feat(ui): image selection gallery state & tweaks

2024-08-30 20:32:17 +00:00 · 2024-05-31 19:20:30 +10:00 · 2024-05-31 19:20:30 +10:00 · e4ce188500
commit e4ce188500
parent e976571fba
19 changed files with 205 additions and 86 deletions
--- a/invokeai/frontend/web/public/locales/en.json
+++ b/invokeai/frontend/web/public/locales/en.json
@ -380,7 +380,11 @@
        "problemDeletingImagesDesc": "One or more images could not be deleted",
        "firstImage": "First Image",
        "secondImage": "Second Image",
-        "selectForCompare": "Select for Compare"
+        "selectForCompare": "Select for Compare",
+        "selectAnImageToCompare": "Select an Image to Compare",
+        "slider": "Slider",
+        "sideBySide": "Side-by-Side",
+        "swapImages": "Swap"
    },
    "hotkeys": {
        "searchHotkeys": "Search Hotkeys",
--- a/invokeai/frontend/web/src/app/components/ThemeLocaleProvider.tsx
+++ b/invokeai/frontend/web/src/app/components/ThemeLocaleProvider.tsx
@ -19,6 +19,13 @@ function ThemeLocaleProvider({ children }: ThemeLocaleProviderProps) {
    return extendTheme({
      ..._theme,
      direction,
+      shadows: {
+        ..._theme.shadows,
+        selectedForCompare:
+          '0px 0px 0px 1px var(--invoke-colors-base-900), 0px 0px 0px 4px var(--invoke-colors-green-400)',
+        hoverSelectedForCompare:
+          '0px 0px 0px 1px var(--invoke-colors-base-900), 0px 0px 0px 4px var(--invoke-colors-green-300)',
+      },
    });
  }, [direction]);

--- a/invokeai/frontend/web/src/common/components/IAIDndImage.tsx
+++ b/invokeai/frontend/web/src/common/components/IAIDndImage.tsx
@ -35,6 +35,7 @@ type IAIDndImageProps = FlexProps & {
  draggableData?: TypesafeDraggableData;
  dropLabel?: ReactNode;
  isSelected?: boolean;
+  isSelectedForCompare?: boolean;
  thumbnail?: boolean;
  noContentFallback?: ReactElement;
  useThumbailFallback?: boolean;
@ -61,6 +62,7 @@ const IAIDndImage = (props: IAIDndImageProps) => {
    draggableData,
    dropLabel,
    isSelected = false,
+    isSelectedForCompare = false,
    thumbnail = false,
    noContentFallback = defaultNoContentFallback,
    uploadElement = defaultUploadElement,
@ -165,7 +167,11 @@ const IAIDndImage = (props: IAIDndImageProps) => {
                data-testid={dataTestId}
              />
              {withMetadataOverlay && <ImageMetadataOverlay imageDTO={imageDTO} />}
-              <SelectionOverlay isSelected={isSelected} isHovered={withHoverOverlay ? isHovered : false} />
+              <SelectionOverlay
+                isSelected={isSelected}
+                isSelectedForCompare={isSelectedForCompare}
+                isHovered={withHoverOverlay ? isHovered : false}
+              />
            </Flex>
          )}
          {!imageDTO && !isUploadDisabled && (
--- a/invokeai/frontend/web/src/common/components/SelectionOverlay.tsx
+++ b/invokeai/frontend/web/src/common/components/SelectionOverlay.tsx
@ -3,10 +3,17 @@ import { memo, useMemo } from 'react';

 type Props = {
  isSelected: boolean;
+  isSelectedForCompare: boolean;
  isHovered: boolean;
 };
-const SelectionOverlay = ({ isSelected, isHovered }: Props) => {
+const SelectionOverlay = ({ isSelected, isSelectedForCompare, isHovered }: Props) => {
  const shadow = useMemo(() => {
+    if (isSelectedForCompare && isHovered) {
+      return 'hoverSelectedForCompare';
+    }
+    if (isSelectedForCompare && !isHovered) {
+      return 'selectedForCompare';
+    }
    if (isSelected && isHovered) {
      return 'hoverSelected';
    }
@ -17,7 +24,7 @@ const SelectionOverlay = ({ isSelected, isHovered }: Props) => {
      return 'hoverUnselected';
    }
    return undefined;
-  }, [isHovered, isSelected]);
+  }, [isHovered, isSelected, isSelectedForCompare]);
  return (
    <Box
      className="selection-box"
@ -27,7 +34,7 @@ const SelectionOverlay = ({ isSelected, isHovered }: Props) => {
      bottom={0}
      insetInlineStart={0}
      borderRadius="base"
-      opacity={isSelected ? 1 : 0.7}
+      opacity={isSelected || isSelectedForCompare ? 1 : 0.7}
      transitionProperty="common"
      transitionDuration="0.1s"
      pointerEvents="none"
--- a/invokeai/frontend/web/src/features/dnd/types/index.ts
+++ b/invokeai/frontend/web/src/features/dnd/types/index.ts
@ -81,6 +81,10 @@ export type RemoveFromBoardDropData = BaseDropData & {

 export type SelectForCompareDropData = BaseDropData & {
  actionType: 'SELECT_FOR_COMPARE';
+  context: {
+    firstImageName?: string | null;
+    secondImageName?: string | null;
+  };
 };

 export type TypesafeDroppableData =
@ -139,7 +143,7 @@ export type UseDraggableTypesafeReturnValue = Omit<ReturnType<typeof useOriginal
  over: TypesafeOver | null;
 };

-export interface TypesafeActive extends Omit<Active, 'data'> {
+interface TypesafeActive extends Omit<Active, 'data'> {
  data: React.MutableRefObject<TypesafeDraggableData | undefined>;
 }

--- a/invokeai/frontend/web/src/features/dnd/util/isValidDrop.ts
+++ b/invokeai/frontend/web/src/features/dnd/util/isValidDrop.ts
@ -33,7 +33,9 @@ export const isValidDrop = (overData?: TypesafeDroppableData | null, activeData?
      return (
        payloadType === 'IMAGE_DTO' &&
        activeData.id !== 'image-compare-first-image' &&
-        activeData.id !== 'image-compare-second-image'
+        activeData.id !== 'image-compare-second-image' &&
+        activeData.payload.imageDTO.image_name !== overData.context.firstImageName &&
+        activeData.payload.imageDTO.image_name !== overData.context.secondImageName
      );
    case 'ADD_TO_BOARD': {
      // If the board is the same, don't allow the drop
--- a/invokeai/frontend/web/src/features/gallery/components/Boards/BoardsList/GalleryBoard.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/Boards/BoardsList/GalleryBoard.tsx
@ -162,7 +162,7 @@ const GalleryBoard = ({ board, isSelected, setBoardToDelete }: GalleryBoardProps
                  </Flex>
                )}
                {isSelectedForAutoAdd && <AutoAddIcon />}
-                <SelectionOverlay isSelected={isSelected} isHovered={isHovered} />
+                <SelectionOverlay isSelected={isSelected} isSelectedForCompare={false} isHovered={isHovered} />
                <Flex
                  position="absolute"
                  bottom={0}
--- a/invokeai/frontend/web/src/features/gallery/components/Boards/BoardsList/NoBoardBoard.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/Boards/BoardsList/NoBoardBoard.tsx
@ -117,7 +117,7 @@ const NoBoardBoard = memo(({ isSelected }: Props) => {
                >
                  {boardName}
                </Flex>
-                <SelectionOverlay isSelected={isSelected} isHovered={isHovered} />
+                <SelectionOverlay isSelected={isSelected} isSelectedForCompare={false} isHovered={isHovered} />
                <IAIDroppable data={droppableData} dropLabel={<Text fontSize="md">{t('unifiedCanvas.move')}</Text>} />
              </Flex>
            </Tooltip>
--- a/invokeai/frontend/web/src/features/gallery/components/ImageContextMenu/SingleSelectionMenuItems.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageContextMenu/SingleSelectionMenuItems.tsx
@ -46,6 +46,11 @@ type SingleSelectionMenuItemsProps = {
 const SingleSelectionMenuItems = (props: SingleSelectionMenuItemsProps) => {
  const { imageDTO } = props;
  const optimalDimension = useAppSelector(selectOptimalDimension);
+  const maySelectForCompare = useAppSelector(
+    (s) =>
+      s.gallery.imageToCompare?.image_name !== imageDTO.image_name &&
+      s.gallery.selection.slice(-1)[0]?.image_name !== imageDTO.image_name
+  );
  const dispatch = useAppDispatch();
  const { t } = useTranslation();
  const isCanvasEnabled = useFeatureStatus('canvas');
@ -136,7 +141,7 @@ const SingleSelectionMenuItems = (props: SingleSelectionMenuItemsProps) => {
      <MenuItem icon={<PiDownloadSimpleBold />} onClickCapture={handleDownloadImage}>
        {t('parameters.downloadImage')}
      </MenuItem>
-      <MenuItem icon={<PiImagesBold />} onClickCapture={handleSelectImageForCompare}>
+      <MenuItem icon={<PiImagesBold />} isDisabled={!maySelectForCompare} onClickCapture={handleSelectImageForCompare}>
        {t('gallery.selectForCompare')}
      </MenuItem>
      <MenuDivider />
--- a/invokeai/frontend/web/src/features/gallery/components/ImageGrid/GalleryImage.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageGrid/GalleryImage.tsx
@ -46,6 +46,9 @@ const GalleryImage = (props: HoverableImageProps) => {
  const { t } = useTranslation();
  const selectedBoardId = useAppSelector((s) => s.gallery.selectedBoardId);
  const alwaysShowImageSizeBadge = useAppSelector((s) => s.gallery.alwaysShowImageSizeBadge);
+  const isSelectedForCompare = useAppSelector(
+    (s) => s.gallery.imageToCompare?.image_name === imageName && s.gallery.viewerMode === 'compare'
+  );
  const { handleClick, isSelected, areMultiplesSelected } = useMultiselect(imageDTO);

  const customStarUi = useStore($customStarUI);
@ -152,6 +155,7 @@ const GalleryImage = (props: HoverableImageProps) => {
          imageDTO={imageDTO}
          draggableData={draggableData}
          isSelected={isSelected}
+          isSelectedForCompare={isSelectedForCompare}
          minSize={0}
          imageSx={imageSx}
          isDropDisabled={true}
--- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentImagePreview.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentImagePreview.tsx
@ -3,10 +3,10 @@ import { createSelector } from '@reduxjs/toolkit';
 import { skipToken } from '@reduxjs/toolkit/query';
 import { useAppSelector } from 'app/store/storeHooks';
 import IAIDndImage from 'common/components/IAIDndImage';
-import IAIDroppable from 'common/components/IAIDroppable';
 import { IAINoContentFallback } from 'common/components/IAIImageFallback';
-import type { SelectForCompareDropData, TypesafeDraggableData } from 'features/dnd/types';
+import type { TypesafeDraggableData } from 'features/dnd/types';
 import ImageMetadataViewer from 'features/gallery/components/ImageMetadataViewer/ImageMetadataViewer';
+import { ImageComparisonDroppable } from 'features/gallery/components/ImageViewer/ImageComparisonDroppable';
 import NextPrevImageButtons from 'features/gallery/components/NextPrevImageButtons';
 import { selectLastSelectedImage } from 'features/gallery/store/gallerySelectors';
 import type { AnimationProps } from 'framer-motion';
@ -23,11 +23,6 @@ const selectLastSelectedImageName = createSelector(
  (lastSelectedImage) => lastSelectedImage?.image_name
 );

-const droppableData: SelectForCompareDropData = {
-  id: 'current-image',
-  actionType: 'SELECT_FOR_COMPARE',
-};
-
 const CurrentImagePreview = () => {
  const { t } = useTranslation();
  const shouldShowImageDetails = useAppSelector((s) => s.ui.shouldShowImageDetails);
@ -85,7 +80,7 @@ const CurrentImagePreview = () => {
          dataTestId="image-preview"
        />
      )}
-      <IAIDroppable data={droppableData} dropLabel="Select for Compare" />
+      <ImageComparisonDroppable />
      {shouldShowImageDetails && imageDTO && (
        <Box position="absolute" opacity={0.8} top={0} width="full" height="full" borderRadius="base">
          <ImageMetadataViewer image={imageDTO} />
--- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageComparison.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageComparison.tsx
@ -1,12 +1,16 @@
 import { createMemoizedSelector } from 'app/store/createMemoizedSelector';
 import { useAppSelector } from 'app/store/storeHooks';
 import IAIDroppable from 'common/components/IAIDroppable';
+import { IAINoContentFallback } from 'common/components/IAIImageFallback';
 import type { SelectForCompareDropData } from 'features/dnd/types';
 import { ImageComparisonSideBySide } from 'features/gallery/components/ImageViewer/ImageComparisonSideBySide';
 import { ImageComparisonSlider } from 'features/gallery/components/ImageViewer/ImageComparisonSlider';
 import { selectGallerySlice } from 'features/gallery/store/gallerySlice';
 import type { PropsWithChildren } from 'react';
-import { memo } from 'react';
+import { memo, useMemo } from 'react';
+import { useTranslation } from 'react-i18next';
+import { PiImagesBold } from 'react-icons/pi';
+import type { ImageDTO } from 'services/api/types';

 const selector = createMemoizedSelector(selectGallerySlice, (gallerySlice) => {
  const firstImage = gallerySlice.selection.slice(-1)[0] ?? null;
@ -15,16 +19,21 @@ const selector = createMemoizedSelector(selectGallerySlice, (gallerySlice) => {
 });

 export const ImageComparison = memo(() => {
+  const { t } = useTranslation();
  const comparisonMode = useAppSelector((s) => s.gallery.comparisonMode);
  const { firstImage, secondImage } = useAppSelector(selector);

  if (!firstImage || !secondImage) {
-    return <ImageComparisonWrapper>Select an image to compare</ImageComparisonWrapper>;
+    return (
+      <ImageComparisonWrapper firstImage={firstImage} secondImage={secondImage}>
+        <IAINoContentFallback label={t('gallery.selectAnImageToCompare')} icon={PiImagesBold} />
+      </ImageComparisonWrapper>
+    );
  }

  if (comparisonMode === 'slider') {
    return (
-      <ImageComparisonWrapper>
+      <ImageComparisonWrapper firstImage={firstImage} secondImage={secondImage}>
        <ImageComparisonSlider firstImage={firstImage} secondImage={secondImage} />
      </ImageComparisonWrapper>
    );
@ -32,7 +41,7 @@ export const ImageComparison = memo(() => {

  if (comparisonMode === 'side-by-side') {
    return (
-      <ImageComparisonWrapper>
+      <ImageComparisonWrapper firstImage={firstImage} secondImage={secondImage}>
        <ImageComparisonSideBySide firstImage={firstImage} secondImage={secondImage} />
      </ImageComparisonWrapper>
    );
@ -41,12 +50,24 @@ export const ImageComparison = memo(() => {

 ImageComparison.displayName = 'ImageComparison';

-const droppableData: SelectForCompareDropData = {
-  id: 'image-comparison',
-  actionType: 'SELECT_FOR_COMPARE',
-};
+type Props = PropsWithChildren<{
+  firstImage: ImageDTO | null;
+  secondImage: ImageDTO | null;
+}>;
+
+const ImageComparisonWrapper = memo((props: Props) => {
+  const droppableData = useMemo<SelectForCompareDropData>(
+    () => ({
+      id: 'image-comparison',
+      actionType: 'SELECT_FOR_COMPARE',
+      context: {
+        firstImageName: props.firstImage?.image_name,
+        secondImageName: props.secondImage?.image_name,
+      },
+    }),
+    [props.firstImage?.image_name, props.secondImage?.image_name]
+  );

-const ImageComparisonWrapper = memo((props: PropsWithChildren) => {
  return (
    <>
      {props.children}
--- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageComparisonDroppable.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageComparisonDroppable.tsx
@ -0,0 +1,33 @@
+import { createMemoizedSelector } from 'app/store/createMemoizedSelector';
+import { useAppSelector } from 'app/store/storeHooks';
+import IAIDroppable from 'common/components/IAIDroppable';
+import type { SelectForCompareDropData } from 'features/dnd/types';
+import { selectGallerySlice } from 'features/gallery/store/gallerySlice';
+import { memo, useMemo } from 'react';
+import { useTranslation } from 'react-i18next';
+
+const selector = createMemoizedSelector(selectGallerySlice, (gallerySlice) => {
+  const firstImage = gallerySlice.selection.slice(-1)[0] ?? null;
+  const secondImage = gallerySlice.imageToCompare;
+  return { firstImage, secondImage };
+});
+
+export const ImageComparisonDroppable = memo(() => {
+  const { t } = useTranslation();
+  const { firstImage, secondImage } = useAppSelector(selector);
+  const droppableData = useMemo<SelectForCompareDropData>(
+    () => ({
+      id: 'image-comparison',
+      actionType: 'SELECT_FOR_COMPARE',
+      context: {
+        firstImageName: firstImage?.image_name,
+        secondImageName: secondImage?.image_name,
+      },
+    }),
+    [firstImage?.image_name, secondImage?.image_name]
+  );
+
+  return <IAIDroppable data={droppableData} dropLabel={t('gallery.selectForCompare')} />;
+});
+
+ImageComparisonDroppable.displayName = 'ImageComparisonDroppable';
--- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageComparisonSideBySide.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageComparisonSideBySide.tsx
@ -51,7 +51,12 @@ export const ImageComparisonSideBySide = memo(({ firstImage, secondImage }: Prop
        <PanelGroup ref={panelGroupRef} direction="horizontal" id="image-comparison-side-by-side">
          <Panel minSize={20}>
            <Flex w="full" h="full" alignItems="center" justifyContent="center">
-              <IAIDndImage imageDTO={firstImage} isDropDisabled={true} draggableData={firstImageDraggableData} />
+              <IAIDndImage
+                imageDTO={firstImage}
+                isDropDisabled={true}
+                draggableData={firstImageDraggableData}
+                useThumbailFallback
+              />
            </Flex>
          </Panel>
          <ResizeHandle
@ -62,7 +67,12 @@ export const ImageComparisonSideBySide = memo(({ firstImage, secondImage }: Prop

          <Panel minSize={20}>
            <Flex w="full" h="full" alignItems="center" justifyContent="center">
-              <IAIDndImage imageDTO={secondImage} isDropDisabled={true} draggableData={secondImageDraggableData} />
+              <IAIDndImage
+                imageDTO={secondImage}
+                isDropDisabled={true}
+                draggableData={secondImageDraggableData}
+                useThumbailFallback
+              />
            </Flex>
          </Panel>
        </PanelGroup>
--- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageComparisonSlider.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageComparisonSlider.tsx
@ -9,7 +9,7 @@ import type { ImageDTO } from 'services/api/types';

 const DROP_SHADOW = 'drop-shadow(0px 0px 4px rgb(0, 0, 0))';
 const INITIAL_POS = '50%';
-const HANDLE_WIDTH = 2;
+const HANDLE_WIDTH = 1;
 const HANDLE_WIDTH_PX = `${HANDLE_WIDTH}px`;
 const HANDLE_HITBOX = 20;
 const HANDLE_HITBOX_PX = `${HANDLE_HITBOX}px`;
@ -37,52 +37,11 @@ export const ImageComparisonSlider = memo(({ firstImage, secondImage }: Props) =
  // If the container size is not provided, use an internal ref and measure - can cause flicker on mount tho
  const containerRef = useRef<HTMLDivElement>(null);
  const [containerSize] = useMeasure(containerRef);
+  const imageContainerRef = useRef<HTMLDivElement>(null);
  // To keep things smooth, we use RAF to update the handle position & gate it to 60fps
  const rafRef = useRef<number | null>(null);
  const lastMoveTimeRef = useRef<number>(0);

-  const updateHandlePos = useCallback(
-    (clientX: number) => {
-      if (!handleRef.current || !containerRef.current) {
-        return;
-      }
-      lastMoveTimeRef.current = performance.now();
-      const { x, width } = containerRef.current.getBoundingClientRect();
-      const rawHandlePos = ((clientX - x) * 100) / width;
-      const handleWidthPct = (HANDLE_WIDTH * 100) / width;
-      const newHandlePos = Math.min(100 - handleWidthPct, Math.max(0, rawHandlePos));
-      setWidth(`${newHandlePos}%`);
-      setLeft(`calc(${newHandlePos}% - ${HANDLE_HITBOX / 2}px)`);
-    },
-    [containerRef]
-  );
-
-  const onMouseMove = useCallback(
-    (e: MouseEvent) => {
-      if (rafRef.current === null && performance.now() > lastMoveTimeRef.current + 1000 / 60) {
-        rafRef.current = window.requestAnimationFrame(() => {
-          updateHandlePos(e.clientX);
-          rafRef.current = null;
-        });
-      }
-    },
-    [updateHandlePos]
-  );
-
-  const onMouseUp = useCallback(() => {
-    window.removeEventListener('mousemove', onMouseMove);
-  }, [onMouseMove]);
-
-  const onMouseDown = useCallback(
-    (e: React.MouseEvent<HTMLDivElement>) => {
-      // Update the handle position immediately on click
-      updateHandlePos(e.clientX);
-      window.addEventListener('mouseup', onMouseUp, { once: true });
-      window.addEventListener('mousemove', onMouseMove);
-    },
-    [onMouseMove, onMouseUp, updateHandlePos]
-  );
-
  const fittedSize = useMemo<Dimensions>(() => {
    // Fit the first image to the container
    if (containerSize.width === 0 || containerSize.height === 0) {
@ -110,6 +69,45 @@ export const ImageComparisonSlider = memo(({ firstImage, secondImage }: Props) =
    return { width, height };
  }, [containerSize, firstImage.height, firstImage.width]);

+  const updateHandlePos = useCallback((clientX: number) => {
+    if (!handleRef.current || !imageContainerRef.current) {
+      return;
+    }
+    lastMoveTimeRef.current = performance.now();
+    const { x, width } = imageContainerRef.current.getBoundingClientRect();
+    const rawHandlePos = ((clientX - x) * 100) / width;
+    const handleWidthPct = (HANDLE_WIDTH * 100) / width;
+    const newHandlePos = Math.min(100 - handleWidthPct, Math.max(0, rawHandlePos));
+    setWidth(`${newHandlePos}%`);
+    setLeft(`calc(${newHandlePos}% - ${HANDLE_HITBOX / 2}px)`);
+  }, []);
+
+  const onMouseMove = useCallback(
+    (e: MouseEvent) => {
+      if (rafRef.current === null && performance.now() > lastMoveTimeRef.current + 1000 / 60) {
+        rafRef.current = window.requestAnimationFrame(() => {
+          updateHandlePos(e.clientX);
+          rafRef.current = null;
+        });
+      }
+    },
+    [updateHandlePos]
+  );
+
+  const onMouseUp = useCallback(() => {
+    window.removeEventListener('mousemove', onMouseMove);
+  }, [onMouseMove]);
+
+  const onMouseDown = useCallback(
+    (e: React.MouseEvent<HTMLDivElement>) => {
+      // Update the handle position immediately on click
+      updateHandlePos(e.clientX);
+      window.addEventListener('mouseup', onMouseUp, { once: true });
+      window.addEventListener('mousemove', onMouseMove);
+    },
+    [onMouseMove, onMouseUp, updateHandlePos]
+  );
+
  useEffect(
    () => () => {
      if (rafRef.current !== null) {
@ -141,6 +139,7 @@ export const ImageComparisonSlider = memo(({ firstImage, secondImage }: Props) =
        justifyContent="center"
      >
        <Box
+          ref={imageContainerRef}
          position="relative"
          id="image-comparison-second-image-container"
          w={fittedSize.width}
--- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageComparisonToolbarButtons.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ImageComparisonToolbarButtons.tsx
@ -1,9 +1,11 @@
 import { Button, ButtonGroup } from '@invoke-ai/ui-library';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
-import { comparisonModeChanged } from 'features/gallery/store/gallerySlice';
+import { comparedImagesSwapped, comparisonModeChanged } from 'features/gallery/store/gallerySlice';
 import { memo, useCallback } from 'react';
+import { useTranslation } from 'react-i18next';

 export const ImageComparisonToolbarButtons = memo(() => {
+  const { t } = useTranslation();
  const dispatch = useAppDispatch();
  const comparisonMode = useAppSelector((s) => s.gallery.comparisonMode);
  const setComparisonModeSlider = useCallback(() => {
@ -12,26 +14,24 @@ export const ImageComparisonToolbarButtons = memo(() => {
  const setComparisonModeSideBySide = useCallback(() => {
    dispatch(comparisonModeChanged('side-by-side'));
  }, [dispatch]);
-  const setComparisonModeOverlay = useCallback(() => {
-    dispatch(comparisonModeChanged('overlay'));
+  const swapImages = useCallback(() => {
+    dispatch(comparedImagesSwapped());
  }, [dispatch]);

  return (
    <>
      <ButtonGroup variant="outline">
        <Button onClick={setComparisonModeSlider} colorScheme={comparisonMode === 'slider' ? 'invokeBlue' : 'base'}>
-          Slider
+          {t('gallery.slider')}
        </Button>
        <Button
          onClick={setComparisonModeSideBySide}
          colorScheme={comparisonMode === 'side-by-side' ? 'invokeBlue' : 'base'}
        >
-          Side-by-Side
-        </Button>
-        <Button onClick={setComparisonModeOverlay} colorScheme={comparisonMode === 'overlay' ? 'invokeBlue' : 'base'}>
-          Overlay
+          {t('gallery.sideBySide')}
        </Button>
      </ButtonGroup>
+      <Button onClick={swapImages}>{t('gallery.swapImages')}</Button>
    </>
  );
 });
--- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ViewerToggleMenu.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/ViewerToggleMenu.tsx
@ -8,8 +8,9 @@ import {
  PopoverContent,
  PopoverTrigger,
  Text,
+  useDisclosure,
 } from '@invoke-ai/ui-library';
-import { useMemo } from 'react';
+import { useCallback, useMemo } from 'react';
 import { useTranslation } from 'react-i18next';
 import { PiCaretDownBold, PiCheckBold, PiEyeBold, PiImagesBold, PiPencilBold } from 'react-icons/pi';

@ -17,6 +18,7 @@ import { useImageViewer } from './useImageViewer';

 export const ViewerToggleMenu = () => {
  const { t } = useTranslation();
+  const { isOpen, onOpen, onClose } = useDisclosure();
  const { viewerMode, openEditor, openViewer, openCompare } = useImageViewer();
  const icon = useMemo(() => {
    if (viewerMode === 'view') {
@ -40,9 +42,21 @@ export const ViewerToggleMenu = () => {
      return t('common.comparing');
    }
  }, [t, viewerMode]);
+  const _openEditor = useCallback(() => {
+    openEditor();
+    onClose();
+  }, [onClose, openEditor]);
+  const _openViewer = useCallback(() => {
+    openViewer();
+    onClose();
+  }, [onClose, openViewer]);
+  const _openCompare = useCallback(() => {
+    openCompare();
+    onClose();
+  }, [onClose, openCompare]);

  return (
-    <Popover isLazy>
+    <Popover isOpen={isOpen} onClose={onClose} onOpen={onOpen}>
      <PopoverTrigger>
        <Button variant="outline" data-testid="toggle-viewer-menu-button">
          <Flex gap={3} w="full" alignItems="center">
@ -56,7 +70,7 @@ export const ViewerToggleMenu = () => {
        <PopoverArrow />
        <PopoverBody>
          <Flex flexDir="column">
-            <Button onClick={openViewer} variant="ghost" h="auto" w="auto" p={2}>
+            <Button onClick={_openViewer} variant="ghost" h="auto" w="auto" p={2}>
              <Flex gap={2} w="full">
                <Icon as={PiCheckBold} visibility={viewerMode === 'view' ? 'visible' : 'hidden'} />
                <Flex flexDir="column" gap={2} alignItems="flex-start">
@ -69,7 +83,7 @@ export const ViewerToggleMenu = () => {
                </Flex>
              </Flex>
            </Button>
-            <Button onClick={openEditor} variant="ghost" h="auto" w="auto" p={2}>
+            <Button onClick={_openEditor} variant="ghost" h="auto" w="auto" p={2}>
              <Flex gap={2} w="full">
                <Icon as={PiCheckBold} visibility={viewerMode === 'edit' ? 'visible' : 'hidden'} />
                <Flex flexDir="column" gap={2} alignItems="flex-start">
@ -82,7 +96,7 @@ export const ViewerToggleMenu = () => {
                </Flex>
              </Flex>
            </Button>
-            <Button onClick={openCompare} variant="ghost" h="auto" w="auto" p={2}>
+            <Button onClick={_openCompare} variant="ghost" h="auto" w="auto" p={2}>
              <Flex gap={2} w="full">
                <Icon as={PiCheckBold} visibility={viewerMode === 'compare' ? 'visible' : 'hidden'} />
                <Flex flexDir="column" gap={2} alignItems="flex-start">
--- a/invokeai/frontend/web/src/features/gallery/store/gallerySlice.ts
+++ b/invokeai/frontend/web/src/features/gallery/store/gallerySlice.ts
@ -90,6 +90,13 @@ export const gallerySlice = createSlice({
    viewerModeChanged: (state, action: PayloadAction<ViewerMode>) => {
      state.viewerMode = action.payload;
    },
+    comparedImagesSwapped: (state) => {
+      if (state.imageToCompare) {
+        const oldSelection = state.selection;
+        state.selection = [state.imageToCompare];
+        state.imageToCompare = oldSelection[0] ?? null;
+      }
+    },
  },
  extraReducers: (builder) => {
    builder.addMatcher(isAnyBoardDeleted, (state, action) => {
@ -130,6 +137,7 @@ export const {
  viewerModeChanged,
  imageToCompareChanged,
  comparisonModeChanged,
+  comparedImagesSwapped,
 } = gallerySlice.actions;

 const isAnyBoardDeleted = isAnyOf(
--- a/invokeai/frontend/web/src/features/gallery/store/types.ts
+++ b/invokeai/frontend/web/src/features/gallery/store/types.ts
@ -7,7 +7,7 @@ export const IMAGE_LIMIT = 20;

 export type GalleryView = 'images' | 'assets';
 export type BoardId = 'none' | (string & Record<never, never>);
-export type ComparisonMode = 'slider' | 'side-by-side' | 'overlay';
+export type ComparisonMode = 'slider' | 'side-by-side';
 export type ViewerMode = 'edit' | 'view' | 'compare';

 export type GalleryState = {