From 3835fd2f727ca9d164dcce7d9b88663f7a4a8be7 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Wed, 30 Jul 2025 21:01:29 +1000
Subject: [PATCH] feat(ui): zhoosh image comparison ui

---
 .../components/ImageViewer/CompareToolbar.tsx |  5 +-
 .../ImageViewer/ImageComparison.tsx           | 75 ++++++++++++++-----
 .../ImageViewer/ImageComparisonHover.tsx      | 12 +--
 .../ImageViewer/ImageComparisonSlider.tsx     | 12 +--
 .../components/ImageViewer/ImageViewer.tsx    | 50 ++++++-------
 .../ImageViewer/ImageViewerPanel.tsx          | 40 +++-------
 .../gallery/components/ImageViewer/common.ts  |  8 +-
 7 files changed, 114 insertions(+), 88 deletions(-)

diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CompareToolbar.tsx b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CompareToolbar.tsx
index d31f568743..6f799db9d2 100644
--- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CompareToolbar.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CompareToolbar.tsx
@@ -60,7 +60,7 @@ export const CompareToolbar = memo(() => {
   useRegisteredHotkeys({ id: 'nextComparisonMode', category: 'viewer', callback: nextMode, dependencies: [nextMode] });
 
   return (
-    <Flex w="full" px={2} gap={2} bg="base.750" borderTopRadius="base" h={12}>
+    <Flex w="full" justifyContent="center" h={8}>
       <Flex flex={1} justifyContent="center">
         <Flex marginInlineEnd="auto" alignItems="center">
           <IconButton
@@ -85,7 +85,7 @@ export const CompareToolbar = memo(() => {
         </Flex>
       </Flex>
       <Flex flex={1} justifyContent="center">
-        <ButtonGroup variant="outline" alignItems="center">
+        <ButtonGroup size="sm" variant="outline" alignItems="center">
           <Button
             flexShrink={0}
             onClick={setComparisonModeSlider}
@@ -117,6 +117,7 @@ export const CompareToolbar = memo(() => {
             </Flex>
           </Tooltip>
           <Button
+            size="sm"
             variant="link"
             alignSelf="stretch"
             px={2}
diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageComparison.tsx b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageComparison.tsx
index ce1c92e1f3..899b082b44 100644
--- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageComparison.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageComparison.tsx
@@ -1,32 +1,36 @@
-import { Box, Flex } from '@invoke-ai/ui-library';
+import { Box, Divider, Flex } from '@invoke-ai/ui-library';
 import { useAppSelector } from 'app/store/storeHooks';
-import type { ComparisonProps } from 'features/gallery/components/ImageViewer/common';
+import { debounce } from 'es-toolkit';
+import type { ComparisonWrapperProps } from 'features/gallery/components/ImageViewer/common';
+import { selectImageToCompare } from 'features/gallery/components/ImageViewer/common';
 import { CompareToolbar } from 'features/gallery/components/ImageViewer/CompareToolbar';
 import { ImageComparisonDroppable } from 'features/gallery/components/ImageViewer/ImageComparisonDroppable';
 import { ImageComparisonHover } from 'features/gallery/components/ImageViewer/ImageComparisonHover';
 import { ImageComparisonSideBySide } from 'features/gallery/components/ImageViewer/ImageComparisonSideBySide';
 import { ImageComparisonSlider } from 'features/gallery/components/ImageViewer/ImageComparisonSlider';
-import { selectComparisonMode } from 'features/gallery/store/gallerySelectors';
-import { memo } from 'react';
-import { useMeasure } from 'react-use';
+import { selectComparisonMode, selectLastSelectedImage } from 'features/gallery/store/gallerySelectors';
+import { memo, useCallback, useLayoutEffect, useRef, useState } from 'react';
+import { useImageDTO } from 'services/api/endpoints/images';
 import type { Equals } from 'tsafe';
 import { assert } from 'tsafe';
 
-const ImageComparisonContent = memo(({ firstImage, secondImage, containerDims }: ComparisonProps) => {
+const ImageComparisonContent = memo(({ firstImage, secondImage, rect }: ComparisonWrapperProps) => {
   const comparisonMode = useAppSelector(selectComparisonMode);
 
+  if (!firstImage || !secondImage) {
+    return null;
+  }
+
   if (comparisonMode === 'slider') {
-    return <ImageComparisonSlider firstImage={firstImage} secondImage={secondImage} containerDims={containerDims} />;
+    return <ImageComparisonSlider firstImage={firstImage} secondImage={secondImage} rect={rect} />;
   }
 
   if (comparisonMode === 'side-by-side') {
-    return (
-      <ImageComparisonSideBySide firstImage={firstImage} secondImage={secondImage} containerDims={containerDims} />
-    );
+    return <ImageComparisonSideBySide firstImage={firstImage} secondImage={secondImage} rect={rect} />;
   }
 
   if (comparisonMode === 'hover') {
-    return <ImageComparisonHover firstImage={firstImage} secondImage={secondImage} containerDims={containerDims} />;
+    return <ImageComparisonHover firstImage={firstImage} secondImage={secondImage} rect={rect} />;
   }
 
   assert<Equals<never, typeof comparisonMode>>(false);
@@ -34,16 +38,51 @@ const ImageComparisonContent = memo(({ firstImage, secondImage, containerDims }:
 
 ImageComparisonContent.displayName = 'ImageComparisonContent';
 
-export const ImageComparison = memo(({ firstImage, secondImage }: Omit<ComparisonProps, 'containerDims'>) => {
-  const [containerRef, containerDims] = useMeasure<HTMLDivElement>();
+export const ImageComparison = memo(() => {
+  const lastSelectedImageName = useAppSelector(selectLastSelectedImage);
+  const lastSelectedImageDTO = useImageDTO(lastSelectedImageName);
+  const comparisonImageName = useAppSelector(selectImageToCompare);
+  const comparisonImageDTO = useImageDTO(comparisonImageName);
+
+  const [rect, setRect] = useState<DOMRect | null>(null);
+  const ref = useRef<HTMLDivElement | null>(null);
+
+  // Ref callback runs synchronously when the DOM node is attached, ensuring we have a measurement before
+  // the comparison content is rendered.
+  const measureNode = useCallback((node: HTMLDivElement) => {
+    if (node) {
+      ref.current = node;
+      const boundingRect = node.getBoundingClientRect();
+      setRect(boundingRect);
+    }
+  }, []);
+
+  useLayoutEffect(() => {
+    const el = ref.current;
+    if (!el) {
+      return;
+    }
+    const measureRect = debounce(() => {
+      const boundingRect = el.getBoundingClientRect();
+      setRect(boundingRect);
+    }, 300);
+    const observer = new ResizeObserver(measureRect);
+    observer.observe(el);
+    return () => {
+      observer.disconnect();
+    };
+  }, []);
 
   return (
-    <Flex flexDir="column" w="full" h="full" position="relative">
+    <Flex flexDir="column" w="full" h="full" overflow="hidden" gap={2} position="relative">
       <CompareToolbar />
-      <Box ref={containerRef} w="full" h="full" p={2} overflow="hidden">
-        <ImageComparisonContent firstImage={firstImage} secondImage={secondImage} containerDims={containerDims} />
-      </Box>
-      <ImageComparisonDroppable />
+      <Divider />
+      <Flex w="full" h="full" position="relative">
+        <Box ref={measureNode} w="full" h="full" overflow="hidden">
+          <ImageComparisonContent firstImage={lastSelectedImageDTO} secondImage={comparisonImageDTO} rect={rect} />
+        </Box>
+        <ImageComparisonDroppable />
+      </Flex>
     </Flex>
   );
 });
diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageComparisonHover.tsx b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageComparisonHover.tsx
index 1e9c78f10b..edfada8bc2 100644
--- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageComparisonHover.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageComparisonHover.tsx
@@ -11,14 +11,16 @@ import { memo, useMemo, useRef } from 'react';
 import type { ComparisonProps } from './common';
 import { fitDimsToContainer, getSecondImageDims } from './common';
 
-export const ImageComparisonHover = memo(({ firstImage, secondImage, containerDims }: ComparisonProps) => {
+export const ImageComparisonHover = memo(({ firstImage, secondImage, rect }: ComparisonProps) => {
   const comparisonFit = useAppSelector(selectComparisonFit);
   const imageContainerRef = useRef<HTMLDivElement>(null);
   const mouseOver = useBoolean(false);
-  const fittedDims = useMemo<Dimensions>(
-    () => fitDimsToContainer(containerDims, firstImage),
-    [containerDims, firstImage]
-  );
+  const fittedDims = useMemo<Dimensions>(() => {
+    if (!rect) {
+      return { width: 0, height: 0 };
+    }
+    return fitDimsToContainer(rect, firstImage);
+  }, [firstImage, rect]);
   const compareImageDims = useMemo<Dimensions>(
     () => getSecondImageDims(comparisonFit, fittedDims, firstImage, secondImage),
     [comparisonFit, fittedDims, firstImage, secondImage]
diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageComparisonSlider.tsx b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageComparisonSlider.tsx
index d35a90570c..1f0d64aafe 100644
--- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageComparisonSlider.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageComparisonSlider.tsx
@@ -19,7 +19,7 @@ const HANDLE_HITBOX_PX = `${HANDLE_HITBOX}px`;
 const HANDLE_INNER_LEFT_PX = `${HANDLE_HITBOX / 2 - HANDLE_WIDTH / 2}px`;
 const HANDLE_LEFT_INITIAL_PX = `calc(${INITIAL_POS} - ${HANDLE_HITBOX / 2}px)`;
 
-export const ImageComparisonSlider = memo(({ firstImage, secondImage, containerDims }: ComparisonProps) => {
+export const ImageComparisonSlider = memo(({ firstImage, secondImage, rect }: ComparisonProps) => {
   const comparisonFit = useAppSelector(selectComparisonFit);
 
   // How far the handle is from the left - this will be a CSS calculation that takes into account the handle width
@@ -33,10 +33,12 @@ export const ImageComparisonSlider = memo(({ firstImage, secondImage, containerD
   const rafRef = useRef<number | null>(null);
   const lastMoveTimeRef = useRef<number>(0);
 
-  const fittedDims = useMemo<Dimensions>(
-    () => fitDimsToContainer(containerDims, firstImage),
-    [containerDims, firstImage]
-  );
+  const fittedDims = useMemo<Dimensions>(() => {
+    if (!rect) {
+      return { width: 0, height: 0 };
+    }
+    return fitDimsToContainer(rect, firstImage);
+  }, [firstImage, rect]);
 
   const compareImageDims = useMemo<Dimensions>(
     () => getSecondImageDims(comparisonFit, fittedDims, firstImage, secondImage),
diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageViewer.tsx b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageViewer.tsx
index 2ff613ef40..ec6ee151b5 100644
--- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageViewer.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageViewer.tsx
@@ -1,42 +1,36 @@
+import { Divider, Flex } from '@invoke-ai/ui-library';
 import { useAppSelector } from 'app/store/storeHooks';
-import { selectImageToCompare } from 'features/gallery/components/ImageViewer/common';
+import { setComparisonImageDndTarget } from 'features/dnd/dnd';
+import { DndDropTarget } from 'features/dnd/DndDropTarget';
 import { CurrentImagePreview } from 'features/gallery/components/ImageViewer/CurrentImagePreview';
-import { ImageComparison } from 'features/gallery/components/ImageViewer/ImageComparison';
 import { selectLastSelectedImage } from 'features/gallery/store/gallerySelectors';
 import { memo } from 'react';
+import { useTranslation } from 'react-i18next';
 import { useImageDTO } from 'services/api/endpoints/images';
 
-// type Props = {
-//   closeButton?: ReactNode;
-// };
+import { ViewerToolbar } from './ViewerToolbar';
 
-// const useFocusRegionOptions = {
-//   focusOnMount: true,
-// };
-
-// const FOCUS_REGION_STYLES: SystemStyleObject = {
-//   display: 'flex',
-//   width: 'full',
-//   height: 'full',
-//   position: 'absolute',
-//   flexDirection: 'column',
-//   inset: 0,
-//   alignItems: 'center',
-//   justifyContent: 'center',
-//   overflow: 'hidden',
-// };
+const dndTargetData = setComparisonImageDndTarget.getData();
 
 export const ImageViewer = memo(() => {
+  const { t } = useTranslation();
+
   const lastSelectedImageName = useAppSelector(selectLastSelectedImage);
   const lastSelectedImageDTO = useImageDTO(lastSelectedImageName);
-  const comparisonImageName = useAppSelector(selectImageToCompare);
-  const comparisonImageDTO = useImageDTO(comparisonImageName);
-
-  if (lastSelectedImageDTO && comparisonImageDTO) {
-    return <ImageComparison firstImage={lastSelectedImageDTO} secondImage={comparisonImageDTO} />;
-  }
-
-  return <CurrentImagePreview imageDTO={lastSelectedImageDTO} />;
+  return (
+    <Flex flexDir="column" w="full" h="full" overflow="hidden" gap={2} position="relative">
+      <ViewerToolbar />
+      <Divider />
+      <Flex w="full" h="full" position="relative">
+        <CurrentImagePreview imageDTO={lastSelectedImageDTO} />
+        <DndDropTarget
+          dndTarget={setComparisonImageDndTarget}
+          dndTargetData={dndTargetData}
+          label={t('gallery.selectForCompare')}
+        />
+      </Flex>
+    </Flex>
+  );
 });
 
 ImageViewer.displayName = 'ImageViewer';
diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageViewerPanel.tsx b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageViewerPanel.tsx
index 6442245557..6205d7a4fc 100644
--- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageViewerPanel.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageViewerPanel.tsx
@@ -1,42 +1,24 @@
-import { Divider, Flex } from '@invoke-ai/ui-library';
+import { createSelector } from '@reduxjs/toolkit';
 import { useAppSelector } from 'app/store/storeHooks';
-import type { SetComparisonImageDndTargetData } from 'features/dnd/dnd';
-import { setComparisonImageDndTarget } from 'features/dnd/dnd';
-import { DndDropTarget } from 'features/dnd/DndDropTarget';
 import { selectImageToCompare, selectLastSelectedImage } from 'features/gallery/store/gallerySelectors';
-import { memo, useMemo } from 'react';
-import { useTranslation } from 'react-i18next';
+import { memo } from 'react';
 
 import { ImageViewerContextProvider } from './context';
+import { ImageComparison } from './ImageComparison';
 import { ImageViewer } from './ImageViewer';
-import { ViewerToolbar } from './ViewerToolbar';
+
+const selectIsComparing = createSelector(
+  [selectLastSelectedImage, selectImageToCompare],
+  (lastSelectedImage, imageToCompare) => !!lastSelectedImage && !!imageToCompare
+);
 
 export const ImageViewerPanel = memo(() => {
-  const { t } = useTranslation();
-  const lastSelectedImage = useAppSelector(selectLastSelectedImage);
-  const imageToCompare = useAppSelector(selectImageToCompare);
-
-  // Only show drop target when we have a selected image but no comparison image yet
-  const shouldShowDropTarget = lastSelectedImage && !imageToCompare;
-
-  const dndTargetData = useMemo<SetComparisonImageDndTargetData>(() => setComparisonImageDndTarget.getData(), []);
+  const isComparing = useAppSelector(selectIsComparing);
 
   return (
     <ImageViewerContextProvider>
-      <Flex flexDir="column" w="full" h="full" overflow="hidden" gap={2} position="relative">
-        <ViewerToolbar />
-        <Divider />
-        <Flex w="full" h="full" position="relative">
-          <ImageViewer />
-          {shouldShowDropTarget && (
-            <DndDropTarget
-              dndTarget={setComparisonImageDndTarget}
-              dndTargetData={dndTargetData}
-              label={t('gallery.selectForCompare')}
-            />
-          )}
-        </Flex>
-      </Flex>
+      {!isComparing && <ImageViewer />}
+      {isComparing && <ImageComparison />}
     </ImageViewerContextProvider>
   );
 });
diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/common.ts b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/common.ts
index 34ed7aec31..dfbf62bbbf 100644
--- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/common.ts
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/common.ts
@@ -7,10 +7,16 @@ import type { ImageDTO } from 'services/api/types';
 
 export const DROP_SHADOW = 'drop-shadow(0px 0px 4px rgb(0, 0, 0)) drop-shadow(0px 0px 4px rgba(0, 0, 0, 0.3))';
 
+export type ComparisonWrapperProps = {
+  firstImage: ImageDTO | null;
+  secondImage: ImageDTO | null;
+  rect: DOMRect | null;
+};
+
 export type ComparisonProps = {
   firstImage: ImageDTO;
   secondImage: ImageDTO;
-  containerDims: Dimensions;
+  rect: DOMRect | null;
 };
 
 export const fitDimsToContainer = (containerDims: Dimensions, imageDims: Dimensions): Dimensions => {