From 1446d3490b0291d2d05785c2bf3c6ef0e7c59d2b Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Fri, 30 May 2025 21:31:43 +1000
Subject: [PATCH] fix(ui): merge refs when forwardingin DndImage

---
 .../web/src/features/dnd/DndImage.tsx         | 125 +++++++++---------
 .../ImageContextMenu/ImageContextMenu.tsx     |   5 +-
 2 files changed, 63 insertions(+), 67 deletions(-)

diff --git a/invokeai/frontend/web/src/features/dnd/DndImage.tsx b/invokeai/frontend/web/src/features/dnd/DndImage.tsx
index ded8ab249e..9009fe440f 100644
--- a/invokeai/frontend/web/src/features/dnd/DndImage.tsx
+++ b/invokeai/frontend/web/src/features/dnd/DndImage.tsx
@@ -9,7 +9,7 @@ import { createSingleImageDragPreview, setSingleImageDragPreview } from 'feature
 import { firefoxDndFix } from 'features/dnd/util';
 import { useImageContextMenu } from 'features/gallery/components/ImageContextMenu/ImageContextMenu';
 import { $imageViewer } from 'features/gallery/components/ImageViewer/useImageViewer';
-import { memo, useEffect, useRef, useState } from 'react';
+import { forwardRef, memo, useEffect, useImperativeHandle, useRef, useState } from 'react';
 import type { ImageDTO } from 'services/api/types';
 
 const sx = {
@@ -23,69 +23,68 @@ const sx = {
   },
 } satisfies SystemStyleObject;
 
-/* eslint-disable-next-line @typescript-eslint/no-namespace */
-export namespace DndImage {
-  export interface Props extends ImageProps {
-    imageDTO: ImageDTO;
-    asThumbnail?: boolean;
-  }
-}
+type Props = {
+  imageDTO: ImageDTO;
+  asThumbnail?: boolean;
+} & ImageProps;
 
-export const DndImage = memo(({ imageDTO, asThumbnail, ...rest }: DndImage.Props) => {
-  const store = useAppStore();
-  const [isDragging, setIsDragging] = useState(false);
-  const ref = useRef<HTMLImageElement>(null);
-  const [dragPreviewState, setDragPreviewState] = useState<DndDragPreviewSingleImageState | null>(null);
+export const DndImage = memo(
+  forwardRef(({ imageDTO, asThumbnail, ...rest }: Props, forwardedRef) => {
+    const store = useAppStore();
+    const [isDragging, setIsDragging] = useState(false);
+    const ref = useRef<HTMLImageElement>(null);
+    useImperativeHandle(forwardedRef, () => ref.current!, []);
+    const [dragPreviewState, setDragPreviewState] = useState<DndDragPreviewSingleImageState | null>(null);
 
-  useEffect(() => {
-    const element = ref.current;
-    if (!element) {
-      return;
-    }
-    return combine(
-      firefoxDndFix(element),
-      draggable({
-        element,
-        getInitialData: () => singleImageDndSource.getData({ imageDTO }, imageDTO.image_name),
-        onDragStart: () => {
-          setIsDragging(true);
-          if ($imageViewer.get()) {
-            $imageViewer.set(false);
-          }
-        },
-        onDrop: () => {
-          setIsDragging(false);
-        },
-        onGenerateDragPreview: (args) => {
-          if (singleImageDndSource.typeGuard(args.source.data)) {
-            setSingleImageDragPreview({
-              singleImageDndData: args.source.data,
-              onGenerateDragPreviewArgs: args,
-              setDragPreviewState,
-            });
-          }
-        },
-      })
+    useEffect(() => {
+      const element = ref.current;
+      if (!element) {
+        return;
+      }
+      return combine(
+        firefoxDndFix(element),
+        draggable({
+          element,
+          getInitialData: () => singleImageDndSource.getData({ imageDTO }, imageDTO.image_name),
+          onDragStart: () => {
+            setIsDragging(true);
+            if ($imageViewer.get()) {
+              $imageViewer.set(false);
+            }
+          },
+          onDrop: () => {
+            setIsDragging(false);
+          },
+          onGenerateDragPreview: (args) => {
+            if (singleImageDndSource.typeGuard(args.source.data)) {
+              setSingleImageDragPreview({
+                singleImageDndData: args.source.data,
+                onGenerateDragPreviewArgs: args,
+                setDragPreviewState,
+              });
+            }
+          },
+        })
+      );
+    }, [forwardedRef, imageDTO, store]);
+
+    useImageContextMenu(imageDTO, ref);
+
+    return (
+      <>
+        <Image
+          role="button"
+          ref={ref}
+          src={asThumbnail ? imageDTO.thumbnail_url : imageDTO.image_url}
+          fallbackSrc={asThumbnail ? undefined : imageDTO.thumbnail_url}
+          w={imageDTO.width}
+          sx={sx}
+          data-is-dragging={isDragging}
+          {...rest}
+        />
+        {dragPreviewState?.type === 'single-image' ? createSingleImageDragPreview(dragPreviewState) : null}
+      </>
     );
-  }, [imageDTO, store]);
-
-  useImageContextMenu(imageDTO, ref);
-
-  return (
-    <>
-      <Image
-        role="button"
-        ref={ref}
-        src={asThumbnail ? imageDTO.thumbnail_url : imageDTO.image_url}
-        fallbackSrc={asThumbnail ? undefined : imageDTO.thumbnail_url}
-        w={imageDTO.width}
-        sx={sx}
-        data-is-dragging={isDragging}
-        {...rest}
-      />
-      {dragPreviewState?.type === 'single-image' ? createSingleImageDragPreview(dragPreviewState) : null}
-    </>
-  );
-});
-
+  })
+);
 DndImage.displayName = 'DndImage';
diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageContextMenu/ImageContextMenu.tsx b/invokeai/frontend/web/src/features/gallery/components/ImageContextMenu/ImageContextMenu.tsx
index b8676a2ad6..d944d3dac0 100644
--- a/invokeai/frontend/web/src/features/gallery/components/ImageContextMenu/ImageContextMenu.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageContextMenu/ImageContextMenu.tsx
@@ -60,11 +60,8 @@ 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, ref: RefObject<HTMLElement>) => {
+export const useImageContextMenu = (imageDTO: ImageDTO, ref: RefObject<HTMLElement>) => {
   useEffect(() => {
-    if (!imageDTO) {
-      return;
-    }
     const el = ref.current;
     if (!el) {
       return;