refactor(ui): image viewer & comparison convolutedness

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

@@ -1,29 +1,23 @@
 import { Box, Flex } from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
-import { skipToken } from '@reduxjs/toolkit/query';
 import { useAppSelector } from 'app/store/storeHooks';
 import { CanvasAlertsInvocationProgress } from 'features/controlLayers/components/CanvasAlerts/CanvasAlertsInvocationProgress';
 import { CanvasAlertsSendingToCanvas } from 'features/controlLayers/components/CanvasAlerts/CanvasAlertsSendingTo';
 import { DndImage } from 'features/dnd/DndImage';
 import ImageMetadataViewer from 'features/gallery/components/ImageMetadataViewer/ImageMetadataViewer';
 import NextPrevImageButtons from 'features/gallery/components/NextPrevImageButtons';
-import { selectLastSelectedImageName } from 'features/gallery/store/gallerySelectors';
 import { selectShouldShowImageDetails, selectShouldShowProgressInViewer } from 'features/ui/store/uiSelectors';
 import type { AnimationProps } from 'framer-motion';
 import { AnimatePresence, motion } from 'framer-motion';
 import { memo, useCallback, useRef, useState } from 'react';
-import { useGetImageDTOQuery } from 'services/api/endpoints/images';
 import type { ImageDTO } from 'services/api/types';
 import { $hasProgressImage, $isProgressFromCanvas } from 'services/events/stores';
 
 import { NoContentForViewer } from './NoContentForViewer';
 import ProgressImage from './ProgressImage';
 
-const CurrentImagePreview = () => {
+const CurrentImagePreview = ({ imageDTO }: { imageDTO?: ImageDTO }) => {
   const shouldShowImageDetails = useAppSelector(selectShouldShowImageDetails);
-  const imageName = useAppSelector(selectLastSelectedImageName);
-
-  const { currentData: imageDTO } = useGetImageDTOQuery(imageName ?? skipToken);
 
   // Show and hide the next/prev buttons on mouse move
   const [shouldShowNextPrevButtons, setShouldShowNextPrevButtons] = useState<boolean>(false);

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

@@ -1,42 +1,50 @@
+import { Box, Flex } from '@invoke-ai/ui-library';
 import { useAppSelector } from 'app/store/storeHooks';
-import { IAINoContentFallback } from 'common/components/IAIImageFallback';
-import type { Dimensions } from 'features/controlLayers/store/types';
-import { selectComparisonImages } from 'features/gallery/components/ImageViewer/common';
+import type { ComparisonProps } 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 { useTranslation } from 'react-i18next';
-import { PiImagesBold } from 'react-icons/pi';
+import { useMeasure } from 'react-use';
+import type { Equals } from 'tsafe';
+import { assert } from 'tsafe';
 
-type Props = {
-  containerDims: Dimensions;
-};
-
-export const ImageComparison = memo(({ containerDims }: Props) => {
-  const { t } = useTranslation();
+export const ImageComparisonContent = memo(({ firstImage, secondImage, containerDims }: ComparisonProps) => {
   const comparisonMode = useAppSelector(selectComparisonMode);
-  const { firstImage, secondImage } = useAppSelector(selectComparisonImages);
-
-  if (!firstImage || !secondImage) {
-    // Should rarely/never happen - we don't render this component unless we have images to compare
-    return <IAINoContentFallback label={t('gallery.selectAnImageToCompare')} icon={PiImagesBold} />;
-  }
 
   if (comparisonMode === 'slider') {
-    return <ImageComparisonSlider containerDims={containerDims} firstImage={firstImage} secondImage={secondImage} />;
+    return <ImageComparisonSlider firstImage={firstImage} secondImage={secondImage} containerDims={containerDims} />;
   }
 
   if (comparisonMode === 'side-by-side') {
     return (
-      <ImageComparisonSideBySide containerDims={containerDims} firstImage={firstImage} secondImage={secondImage} />
+      <ImageComparisonSideBySide firstImage={firstImage} secondImage={secondImage} containerDims={containerDims} />
     );
   }
 
   if (comparisonMode === 'hover') {
-    return <ImageComparisonHover containerDims={containerDims} firstImage={firstImage} secondImage={secondImage} />;
+    return <ImageComparisonHover firstImage={firstImage} secondImage={secondImage} containerDims={containerDims} />;
   }
+
+  assert<Equals<never, typeof comparisonMode>>(false);
 });
 
+ImageComparisonContent.displayName = 'ImageComparisonContent';
+
+export const ImageComparison = memo(({ firstImage, secondImage }: Omit<ComparisonProps, 'containerDims'>) => {
+  const [containerRef, containerDims] = useMeasure<HTMLDivElement>();
+
+  return (
+    <Flex flexDir="column" w="full" h="full" position="relative">
+      <CompareToolbar />
+      <Box ref={containerRef} w="full" h="full" p={2} overflow="hidden">
+        <ImageComparisonContent firstImage={firstImage} secondImage={secondImage} containerDims={containerDims} />
+      </Box>
+      <ImageComparisonDroppable />
+    </Flex>
+  );
+});
 ImageComparison.displayName = 'ImageComparison';

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

@@ -5,6 +5,7 @@ import { VerticalResizeHandle } from 'features/ui/components/tabs/ResizeHandle';
 import { memo, useCallback, useRef } from 'react';
 import type { ImperativePanelGroupHandle } from 'react-resizable-panels';
 import { Panel, PanelGroup } from 'react-resizable-panels';
+import type { ImageDTO } from 'services/api/types';
 
 export const ImageComparisonSideBySide = memo(({ firstImage, secondImage }: ComparisonProps) => {
   const panelGroupRef = useRef<ImperativePanelGroupHandle>(null);
@@ -25,42 +26,11 @@ export const ImageComparisonSideBySide = memo(({ firstImage, secondImage }: Comp
           autoSaveId="image-comparison-side-by-side"
         >
           <Panel minSize={20}>
-            <Flex position="relative" w="full" h="full" alignItems="center" justifyContent="center">
-              <Flex position="absolute" maxW="full" maxH="full" aspectRatio={firstImage.width / firstImage.height}>
-                <Image
-                  id="image-comparison-side-by-side-first-image"
-                  w={firstImage.width}
-                  h={firstImage.height}
-                  maxW="full"
-                  maxH="full"
-                  src={firstImage.image_url}
-                  fallbackSrc={firstImage.thumbnail_url}
-                  objectFit="contain"
-                  borderRadius="base"
-                />
-                <ImageComparisonLabel type="first" />
-              </Flex>
-            </Flex>
+            <SideBySideImage imageDTO={firstImage} type="first" />
           </Panel>
           <VerticalResizeHandle id="image-comparison-side-by-side-handle" onDoubleClick={onDoubleClickHandle} />
-
           <Panel minSize={20}>
-            <Flex position="relative" w="full" h="full" alignItems="center" justifyContent="center">
-              <Flex position="absolute" maxW="full" maxH="full" aspectRatio={secondImage.width / secondImage.height}>
-                <Image
-                  id="image-comparison-side-by-side-first-image"
-                  w={secondImage.width}
-                  h={secondImage.height}
-                  maxW="full"
-                  maxH="full"
-                  src={secondImage.image_url}
-                  fallbackSrc={secondImage.thumbnail_url}
-                  objectFit="contain"
-                  borderRadius="base"
-                />
-                <ImageComparisonLabel type="second" />
-              </Flex>
-            </Flex>
+            <SideBySideImage imageDTO={secondImage} type="second" />
           </Panel>
         </PanelGroup>
       </Flex>
@@ -69,3 +39,25 @@ export const ImageComparisonSideBySide = memo(({ firstImage, secondImage }: Comp
 });
 
 ImageComparisonSideBySide.displayName = 'ImageComparisonSideBySide';
+
+const SideBySideImage = memo(({ imageDTO, type }: { imageDTO: ImageDTO; type: 'first' | 'second' }) => {
+  return (
+    <Flex position="relative" w="full" h="full" alignItems="center" justifyContent="center">
+      <Flex position="absolute" maxW="full" maxH="full" aspectRatio={imageDTO.width / imageDTO.height}>
+        <Image
+          id={`image-comparison-side-by-side-${type}-image`}
+          w={imageDTO.width}
+          h={imageDTO.height}
+          maxW="full"
+          maxH="full"
+          src={imageDTO.image_url}
+          fallbackSrc={imageDTO.thumbnail_url}
+          objectFit="contain"
+          borderRadius="base"
+        />
+        <ImageComparisonLabel type={type} />
+      </Flex>
+    </Flex>
+  );
+});
+SideBySideImage.displayName = 'SideBySideImage';

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

@@ -21,6 +21,7 @@ const HANDLE_LEFT_INITIAL_PX = `calc(${INITIAL_POS} - ${HANDLE_HITBOX / 2}px)`;
 
 export const ImageComparisonSlider = memo(({ firstImage, secondImage, containerDims }: 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
   const [left, setLeft] = useState(HANDLE_LEFT_INITIAL_PX);
   // How wide the first image is

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

@@ -1,19 +1,18 @@
-import { Box, IconButton, type SystemStyleObject, useOutsideClick } from '@invoke-ai/ui-library';
+import { Box, Flex, IconButton, type SystemStyleObject, useOutsideClick } from '@invoke-ai/ui-library';
+import { skipToken } from '@reduxjs/toolkit/query';
 import { useAppSelector } from 'app/store/storeHooks';
-import { FocusRegionWrapper } from 'common/components/FocusRegionWrapper';
 import { useAssertSingleton } from 'common/hooks/useAssertSingleton';
-import { CompareToolbar } from 'features/gallery/components/ImageViewer/CompareToolbar';
+import { selectImageToCompare } from 'features/gallery/components/ImageViewer/common';
 import CurrentImagePreview from 'features/gallery/components/ImageViewer/CurrentImagePreview';
 import { ImageComparison } from 'features/gallery/components/ImageViewer/ImageComparison';
-import { ImageComparisonDroppable } from 'features/gallery/components/ImageViewer/ImageComparisonDroppable';
 import { ViewerToolbar } from 'features/gallery/components/ImageViewer/ViewerToolbar';
-import { selectHasImageToCompare } from 'features/gallery/store/gallerySelectors';
+import { selectLastSelectedImageName } from 'features/gallery/store/gallerySelectors';
 import type { ReactNode } from 'react';
 import { memo, useRef } from 'react';
 import { useHotkeys } from 'react-hotkeys-hook';
 import { useTranslation } from 'react-i18next';
 import { PiXBold } from 'react-icons/pi';
-import { useMeasure } from 'react-use';
+import { useGetImageDTOQuery } from 'services/api/endpoints/images';
 
 import { useImageViewer } from './useImageViewer';
 
@@ -37,22 +36,16 @@ const FOCUS_REGION_STYLES: SystemStyleObject = {
   overflow: 'hidden',
 };
 
-export const ImageViewer = memo(({ closeButton }: Props) => {
-  useAssertSingleton('ImageViewer');
-  const hasImageToCompare = useAppSelector(selectHasImageToCompare);
-  const [containerRef, containerDims] = useMeasure<HTMLDivElement>();
+export const ImageViewer = memo(() => {
+  const lastSelectedImageName = useAppSelector(selectLastSelectedImageName);
+  const { data: lastSelectedImageDTO } = useGetImageDTOQuery(lastSelectedImageName ?? skipToken);
+  const comparisonImageDTO = useAppSelector(selectImageToCompare);
 
-  return (
-    <FocusRegionWrapper region="viewer" sx={FOCUS_REGION_STYLES} {...useFocusRegionOptions}>
-      {hasImageToCompare && <CompareToolbar />}
-      {!hasImageToCompare && <ViewerToolbar closeButton={closeButton} />}
-      <Box ref={containerRef} w="full" h="full" p={2} overflow="hidden">
-        {!hasImageToCompare && <CurrentImagePreview />}
-        {hasImageToCompare && <ImageComparison containerDims={containerDims} />}
-      </Box>
-      <ImageComparisonDroppable />
-    </FocusRegionWrapper>
-  );
+  if (lastSelectedImageDTO && comparisonImageDTO) {
+    return <ImageComparison firstImage={lastSelectedImageDTO} secondImage={comparisonImageDTO} />;
+  }
+
+  return <CurrentImagePreview imageDTO={lastSelectedImageDTO} />;
 });
 
 ImageViewer.displayName = 'ImageViewer';
@@ -73,16 +66,6 @@ const imageViewerContainerSx: SystemStyleObject = {
   backdropFilter: 'blur(10px) brightness(70%)',
 };
 
-const imageViewerModalSx: SystemStyleObject = {
-  position: 'absolute',
-  bg: 'base.800',
-  borderRadius: 'base',
-  top: 16,
-  right: 16,
-  bottom: 16,
-  left: 16,
-};
-
 export const ImageViewerModal = memo(() => {
   const ref = useRef<HTMLDivElement>(null);
   const imageViewer = useImageViewer();
@@ -93,9 +76,20 @@ export const ImageViewerModal = memo(() => {
 
   return (
     <Box sx={imageViewerContainerSx} data-hidden={!imageViewer.isOpen}>
-      <Box ref={ref} sx={imageViewerModalSx}>
-        <ImageViewer closeButton={<ImageViewerCloseButton />} />
-      </Box>
+      <Flex
+        ref={ref}
+        flexDir="column"
+        position="absolute"
+        bg="base.900"
+        borderRadius="base"
+        top={16}
+        right={16}
+        bottom={16}
+        left={16}
+      >
+        <ViewerToolbar />
+        <ImageViewer />
+      </Flex>
     </Box>
   );
 });

invokeai/frontend/web/src/features/gallery/components/ImageViewer/common.ts

@@ -1,3 +1,4 @@
+import { createSelector } from '@reduxjs/toolkit';
 import { createMemoizedSelector } from 'app/store/createMemoizedSelector';
 import type { Dimensions } from 'features/controlLayers/store/types';
 import { selectGallerySlice } from 'features/gallery/store/gallerySlice';
@@ -62,3 +63,8 @@ export const selectComparisonImages = createMemoizedSelector(selectGallerySlice,
   const secondImage = gallerySlice.imageToCompare;
   return { firstImage, secondImage };
 });
+export const selectFirstImage = createSelector(
+  selectGallerySlice,
+  (gallerySlice) => gallerySlice.selection.slice(-1)[0] ?? null
+);
+export const selectImageToCompare = createSelector(selectGallerySlice, (gallerySlice) => gallerySlice.imageToCompare);