Add drag-and-drop reordering for reference images

invokeai/frontend/web/src/features/controlLayers/components/RefImage/RefImageList.tsx

@@ -1,23 +1,32 @@
-import { Button, Collapse, Divider, Flex, IconButton } from '@invoke-ai/ui-library';
-import { useAppSelector, useAppStore } from 'app/store/storeHooks';
+import { monitorForElements } from '@atlaskit/pragmatic-drag-and-drop/element/adapter';
+import { extractClosestEdge } from '@atlaskit/pragmatic-drag-and-drop-hitbox/closest-edge';
+import { reorderWithEdge } from '@atlaskit/pragmatic-drag-and-drop-hitbox/util/reorder-with-edge';
+import { Box, Button, Collapse, Divider, Flex, IconButton } from '@invoke-ai/ui-library';
+import { useAppDispatch, useAppSelector, useAppStore } from 'app/store/storeHooks';
 import { useImageUploadButton } from 'common/hooks/useImageUploadButton';
+import { colorTokenToCssVar } from 'common/util/colorTokenToCssVar';
 import { RefImagePreview } from 'features/controlLayers/components/RefImage/RefImagePreview';
+import { useRefImageListDnd } from 'features/controlLayers/components/RefImage/useRefImageListDnd';
 import { CanvasManagerProviderGate } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
-import { RefImageIdContext } from 'features/controlLayers/contexts/RefImageIdContext';
+import { RefImageIdContext, useRefImageIdContext } from 'features/controlLayers/contexts/RefImageIdContext';
 import { getDefaultRefImageConfig } from 'features/controlLayers/hooks/addLayerHooks';
 import { useNewGlobalReferenceImageFromBbox } from 'features/controlLayers/hooks/saveCanvasHooks';
 import { useCanvasIsBusySafe } from 'features/controlLayers/hooks/useCanvasIsBusy';
 import {
   refImageAdded,
+  refImagesReordered,
   selectIsRefImagePanelOpen,
   selectRefImageEntityIds,
   selectSelectedRefEntityId,
 } from 'features/controlLayers/store/refImagesSlice';
 import { imageDTOToCroppableImage } from 'features/controlLayers/store/util';
-import { addGlobalReferenceImageDndTarget } from 'features/dnd/dnd';
+import { addGlobalReferenceImageDndTarget, singleRefImageDndSource } from 'features/dnd/dnd';
 import { DndDropTarget } from 'features/dnd/DndDropTarget';
+import { DndListDropIndicator } from 'features/dnd/DndListDropIndicator';
+import { triggerPostMoveFlash } from 'features/dnd/util';
 import { selectActiveTab } from 'features/ui/store/uiSelectors';
-import { memo, useMemo } from 'react';
+import { memo, useEffect, useMemo, useRef } from 'react';
+import { flushSync } from 'react-dom';
 import { useTranslation } from 'react-i18next';
 import { PiBoundingBoxBold, PiUploadBold } from 'react-icons/pi';
 import type { ImageDTO } from 'services/api/types';
@@ -26,16 +35,79 @@ import { RefImageHeader } from './RefImageHeader';
 import { RefImageSettings } from './RefImageSettings';
 
 export const RefImageList = memo(() => {
+  const dispatch = useAppDispatch();
   const ids = useAppSelector(selectRefImageEntityIds);
   const isPanelOpen = useAppSelector(selectIsRefImagePanelOpen);
   const selectedEntityId = useAppSelector(selectSelectedRefEntityId);
 
+  useEffect(() => {
+    return monitorForElements({
+      canMonitor({ source }) {
+        return singleRefImageDndSource.typeGuard(source.data);
+      },
+      onDrop({ location, source }) {
+        const target = location.current.dropTargets[0];
+        if (!target) {
+          return;
+        }
+
+        const sourceData = source.data;
+        const targetData = target.data;
+
+        if (!singleRefImageDndSource.typeGuard(sourceData) || !singleRefImageDndSource.typeGuard(targetData)) {
+          return;
+        }
+
+        const indexOfSource = ids.indexOf(sourceData.payload.id);
+        const indexOfTarget = ids.indexOf(targetData.payload.id);
+
+        if (indexOfTarget < 0 || indexOfSource < 0) {
+          return;
+        }
+
+        if (indexOfSource === indexOfTarget) {
+          return;
+        }
+
+        const closestEdgeOfTarget = extractClosestEdge(targetData);
+
+        let edgeIndexDelta = 0;
+        if (closestEdgeOfTarget === 'right') {
+          edgeIndexDelta = 1;
+        } else if (closestEdgeOfTarget === 'left') {
+          edgeIndexDelta = -1;
+        }
+
+        if (indexOfSource === indexOfTarget + edgeIndexDelta) {
+          return;
+        }
+
+        const reorderedIds = reorderWithEdge({
+          list: ids,
+          startIndex: indexOfSource,
+          indexOfTarget,
+          closestEdgeOfTarget,
+          axis: 'horizontal',
+        });
+
+        flushSync(() => {
+          dispatch(refImagesReordered({ ids: reorderedIds }));
+        });
+
+        const element = document.querySelector(`[data-ref-image-id="${sourceData.payload.id}"]`);
+        if (element instanceof HTMLElement) {
+          triggerPostMoveFlash(element, colorTokenToCssVar('base.700'));
+        }
+      },
+    });
+  }, [dispatch, ids]);
+
   return (
     <Flex flexDir="column">
       <Flex gap={2} h={16}>
         {ids.map((id) => (
           <RefImageIdContext.Provider key={id} value={id}>
-            <RefImagePreview />
+            <RefImageListItem />
           </RefImageIdContext.Provider>
         ))}
         {ids.length < 5 && <AddRefImageDropTargetAndButton />}
@@ -60,6 +132,20 @@ export const RefImageList = memo(() => {
 
 RefImageList.displayName = 'RefImageList';
 
+const RefImageListItem = memo(() => {
+  const id = useRefImageIdContext();
+  const ref = useRef<HTMLDivElement>(null);
+  const [dndListState, isDragging] = useRefImageListDnd(ref, id);
+
+  return (
+    <Box ref={ref} position="relative" h="full" data-ref-image-id={id} opacity={isDragging ? 0.3 : 1} cursor="grab">
+      <RefImagePreview />
+      <DndListDropIndicator dndState={dndListState} gap="var(--invoke-space-2)" />
+    </Box>
+  );
+});
+RefImageListItem.displayName = 'RefImageListItem';
+
 const dndTargetData = addGlobalReferenceImageDndTarget.getData();
 
 const MaxRefImages = memo(() => {

invokeai/frontend/web/src/features/controlLayers/components/RefImage/useRefImageListDnd.ts

@@ -0,0 +1,81 @@
+import { combine } from '@atlaskit/pragmatic-drag-and-drop/combine';
+import { draggable, dropTargetForElements } from '@atlaskit/pragmatic-drag-and-drop/element/adapter';
+import { attachClosestEdge, extractClosestEdge } from '@atlaskit/pragmatic-drag-and-drop-hitbox/closest-edge';
+import { singleRefImageDndSource } from 'features/dnd/dnd';
+import { type DndListTargetState, idle } from 'features/dnd/types';
+import { firefoxDndFix } from 'features/dnd/util';
+import type { RefObject } from 'react';
+import { useEffect, useState } from 'react';
+
+export const useRefImageListDnd = (ref: RefObject<HTMLElement>, id: string) => {
+  const [dndListState, setDndListState] = useState<DndListTargetState>(idle);
+  const [isDragging, setIsDragging] = useState(false);
+
+  useEffect(() => {
+    const element = ref.current;
+    if (!element) {
+      return;
+    }
+    return combine(
+      firefoxDndFix(element),
+      draggable({
+        element,
+        getInitialData() {
+          return singleRefImageDndSource.getData({ id });
+        },
+        onDragStart() {
+          setDndListState({ type: 'is-dragging' });
+          setIsDragging(true);
+        },
+        onDrop() {
+          setDndListState(idle);
+          setIsDragging(false);
+        },
+      }),
+      dropTargetForElements({
+        element,
+        canDrop({ source }) {
+          if (!singleRefImageDndSource.typeGuard(source.data)) {
+            return false;
+          }
+          return true;
+        },
+        getData({ input }) {
+          const data = singleRefImageDndSource.getData({ id });
+          return attachClosestEdge(data, {
+            element,
+            input,
+            allowedEdges: ['left', 'right'],
+          });
+        },
+        getIsSticky() {
+          return true;
+        },
+        onDragEnter({ self }) {
+          const closestEdge = extractClosestEdge(self.data);
+          setDndListState({ type: 'is-dragging-over', closestEdge });
+        },
+        onDrag({ self }) {
+          const closestEdge = extractClosestEdge(self.data);
+
+          // Only need to update react state if nothing has changed.
+          // Prevents re-rendering.
+          setDndListState((current) => {
+            if (current.type === 'is-dragging-over' && current.closestEdge === closestEdge) {
+              return current;
+            }
+            return { type: 'is-dragging-over', closestEdge };
+          });
+        },
+        onDragLeave() {
+          setDndListState(idle);
+        },
+        onDrop() {
+          setDndListState(idle);
+        },
+      })
+    );
+  }, [id, ref]);
+
+  return [dndListState, isDragging] as const;
+};

invokeai/frontend/web/src/features/controlLayers/store/refImagesSlice.ts

@@ -245,6 +245,17 @@ const slice = createSlice({
       // Preserve the existing image when replacing the config
       entity.config = { ...config, image: entity.config.image };
     },
+    refImagesReordered: (state, action: PayloadAction<{ ids: string[] }>) => {
+      const { ids } = action.payload;
+      const sortedEntities: RefImageState[] = [];
+      for (const id of ids) {
+        const entity = state.entities.find((e) => e.id === id);
+        if (entity) {
+          sortedEntities.push(entity);
+        }
+      }
+      state.entities = sortedEntities;
+    },
     refImagesReset: () => getInitialRefImagesState(),
   },
 });
@@ -263,6 +274,7 @@ export const {
   refImageFLUXReduxImageInfluenceChanged,
   refImageIsEnabledToggled,
   refImagesRecalled,
+  refImagesReordered,
 } = slice.actions;
 
 export const refImagesSliceConfig: SliceConfig<typeof slice> = {

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

@@ -109,6 +109,16 @@ export const singleCanvasEntityDndSource: DndSource<SingleCanvasEntityDndSourceD
   getData: buildGetData(_singleCanvasEntity.key, _singleCanvasEntity.type),
 };
 
+//#region Single Ref Image
+const _singleRefImage = buildTypeAndKey('single-ref-image');
+type SingleRefImageDndSourceData = DndData<typeof _singleRefImage.type, typeof _singleRefImage.key, { id: string }>;
+export const singleRefImageDndSource: DndSource<SingleRefImageDndSourceData> = {
+  ..._singleRefImage,
+  typeGuard: buildTypeGuard(_singleRefImage.key),
+  getData: buildGetData(_singleRefImage.key, _singleRefImage.type),
+};
+//#endregion
+
 type DndTarget<TargetData extends DndData, SourceData extends DndData> = {
   key: symbol;
   type: TargetData['type'];