Fix progress preview gallery selection (invoke-ai#9217)

* Fix progress preview gallery selection

* Remove redundant clearTimeout in unmount cleanup effect

The reveal effect's own cleanup already clears selectedImageRevealTimeoutId
on unmount; the unmount-only effect just needs to reset the atom.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: 3 people

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

@@ -30,10 +30,15 @@ export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | nu
   const shouldShowItemDetails = useAppSelector(selectShouldShowItemDetails);
   const shouldShowProgressInViewer = useAppSelector(selectShouldShowProgressInViewer);
   const { goToPreviousImage, goToNextImage, isFetching } = useNextPrevItemNavigation();
-  const { onLoadImage, $progressEvent, $progressImage } = useImageViewerContext();
+  const { onLoadImage, $progressEvent, $progressImage, $isProgressImageResolving, $isTemporarilyShowingSelectedImage } =
+    useImageViewerContext();
   const progressEvent = useStore($progressEvent);
   const progressImage = useStore($progressImage);
+  const isProgressImageResolving = useStore($isProgressImageResolving);
+  const isTemporarilyShowingSelectedImage = useStore($isTemporarilyShowingSelectedImage);
   const [imageToRender, setImageToRender] = useState<ImageDTO | null>(null);
+  const previousRenderedImageNameRef = useRef<string | null>(null);
+  const selectedImageRevealTimeoutId = useRef(0);
 
   useEffect(() => {
     if (!selectedImageName) {
@@ -76,6 +81,53 @@ export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | nu
     };
   }, [imageDTO, imageToRender?.image_name, selectedImageName]);
 
+  const hasProgressImage = progressImage !== null;
+
+  useEffect(() => {
+    const renderedImageName = imageToRender?.image_name ?? null;
+    const previousRenderedImageName = previousRenderedImageNameRef.current;
+    previousRenderedImageNameRef.current = renderedImageName;
+
+    window.clearTimeout(selectedImageRevealTimeoutId.current);
+
+    if (
+      !shouldShowProgressInViewer ||
+      !hasProgressImage ||
+      isProgressImageResolving ||
+      !renderedImageName ||
+      renderedImageName !== selectedImageName
+    ) {
+      $isTemporarilyShowingSelectedImage.set(false);
+      return;
+    }
+
+    if (previousRenderedImageName === null || previousRenderedImageName === renderedImageName) {
+      return;
+    }
+
+    $isTemporarilyShowingSelectedImage.set(true);
+    selectedImageRevealTimeoutId.current = window.setTimeout(() => {
+      $isTemporarilyShowingSelectedImage.set(false);
+    }, SELECTED_IMAGE_REVEAL_DURATION_MS);
+
+    return () => {
+      window.clearTimeout(selectedImageRevealTimeoutId.current);
+    };
+  }, [
+    $isTemporarilyShowingSelectedImage,
+    hasProgressImage,
+    imageToRender?.image_name,
+    isProgressImageResolving,
+    selectedImageName,
+    shouldShowProgressInViewer,
+  ]);
+
+  useEffect(() => {
+    return () => {
+      $isTemporarilyShowingSelectedImage.set(false);
+    };
+  }, [$isTemporarilyShowingSelectedImage]);
+
   // Show and hide the next/prev buttons on mouse move
   const [shouldShowNextPrevButtons, setShouldShowNextPrevButtons] = useState<boolean>(false);
   const timeoutId = useRef(0);
@@ -133,7 +185,7 @@ export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | nu
     dependencies: [onHotkeyNextImage],
   });
 
-  const withProgress = shouldShowProgressInViewer && progressImage !== null;
+  const withProgress = shouldShowProgressInViewer && hasProgressImage && !isTemporarilyShowingSelectedImage;
 
   return (
     <Flex
@@ -202,3 +254,5 @@ const exit: AnimationProps['exit'] = {
   opacity: 0,
   transition: { duration: 0.07 },
 };
+
+const SELECTED_IMAGE_REVEAL_DURATION_MS = 2000;

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

@@ -15,9 +15,10 @@ export const ToggleMetadataViewerButton = memo(() => {
   const dispatch = useAppDispatch();
   const ctx = useImageViewerContext();
   const hasProgressImage = useStore(ctx.$hasProgressImage);
+  const isTemporarilyShowingSelectedImage = useStore(ctx.$isTemporarilyShowingSelectedImage);
   const shouldShowProgressInViewer = useAppSelector(selectShouldShowProgressInViewer);
 
-  const isDisabledOverride = hasProgressImage && shouldShowProgressInViewer;
+  const isDisabledOverride = hasProgressImage && shouldShowProgressInViewer && !isTemporarilyShowingSelectedImage;
 
   const shouldShowItemDetails = useAppSelector(selectShouldShowItemDetails);
   const imageDTO = useAppSelector(selectLastSelectedItem);

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

@@ -4,9 +4,9 @@ import { useAppSelector } from 'app/store/storeHooks';
 import { selectAutoSwitch } from 'features/gallery/store/gallerySelectors';
 import type { ProgressImage as ProgressImageType } from 'features/nodes/types/common';
 import { LRUCache } from 'lru-cache';
-import { type Atom, atom, computed } from 'nanostores';
+import { type Atom, atom, computed, type WritableAtom } from 'nanostores';
 import type { PropsWithChildren } from 'react';
-import { createContext, memo, useCallback, useContext, useEffect, useMemo, useState } from 'react';
+import { createContext, memo, useCallback, useContext, useEffect, useMemo, useRef, useState } from 'react';
 import type { S } from 'services/api/types';
 import { $socket } from 'services/events/stores';
 import { assert } from 'tsafe';
@@ -16,6 +16,8 @@ type ImageViewerContextValue = {
   $progressEvent: Atom<S['InvocationProgressEvent'] | null>;
   $progressImage: Atom<ProgressImageType | null>;
   $hasProgressImage: Atom<boolean>;
+  $isProgressImageResolving: Atom<boolean>;
+  $isTemporarilyShowingSelectedImage: WritableAtom<boolean>;
   onLoadImage: () => void;
 };
 
@@ -29,6 +31,9 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => {
   const $progressEvent = useState(() => atom<S['InvocationProgressEvent'] | null>(null))[0];
   const $progressImage = useState(() => atom<ProgressImageType | null>(null))[0];
   const $hasProgressImage = useState(() => computed($progressImage, (progressImage) => progressImage !== null))[0];
+  const $isProgressImageResolving = useState(() => atom(false))[0];
+  const $isTemporarilyShowingSelectedImage = useState(() => atom(false))[0];
+  const shouldClearProgressImageOnLoadRef = useRef(false);
   // We can have race conditions where we receive a progress event for a queue item that has already finished. Easiest
   // way to handle this is to keep track of finished queue items in a cache and ignore progress events for those.
   const [finishedQueueItemIds] = useState(() => new LRUCache<number, boolean>({ max: 200 }));
@@ -46,6 +51,8 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => {
         );
         return;
       }
+      shouldClearProgressImageOnLoadRef.current = false;
+      $isProgressImageResolving.set(false);
       $progressEvent.set(data);
       if (data.image) {
         $progressImage.set(data.image);
@@ -57,7 +64,7 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => {
     return () => {
       socket.off('invocation_progress', onInvocationProgress);
     };
-  }, [$progressEvent, $progressImage, finishedQueueItemIds, socket]);
+  }, [$isProgressImageResolving, $progressEvent, $progressImage, finishedQueueItemIds, socket]);
 
   useEffect(() => {
     if (!socket) {
@@ -92,8 +99,13 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => {
           // will be stuck on the viewer.
           (data.origin === 'canvas' && data.destination !== 'canvas')
         ) {
+          shouldClearProgressImageOnLoadRef.current = false;
+          $isProgressImageResolving.set(false);
           $progressEvent.set(null);
           $progressImage.set(null);
+        } else {
+          shouldClearProgressImageOnLoadRef.current = true;
+          $isProgressImageResolving.set(true);
         }
       }
     };
@@ -103,16 +115,36 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => {
     return () => {
       socket.off('queue_item_status_changed', onQueueItemStatusChanged);
     };
-  }, [$progressEvent, $progressImage, autoSwitch, finishedQueueItemIds, socket]);
+  }, [$isProgressImageResolving, $progressEvent, $progressImage, autoSwitch, finishedQueueItemIds, socket]);
 
   const onLoadImage = useCallback(() => {
+    if (!shouldClearProgressImageOnLoadRef.current) {
+      return;
+    }
+
+    shouldClearProgressImageOnLoadRef.current = false;
+    $isProgressImageResolving.set(false);
     $progressEvent.set(null);
     $progressImage.set(null);
-  }, [$progressEvent, $progressImage]);
+  }, [$isProgressImageResolving, $progressEvent, $progressImage]);
 
   const value = useMemo(
-    () => ({ $progressEvent, $progressImage, $hasProgressImage, onLoadImage }),
-    [$hasProgressImage, $progressEvent, $progressImage, onLoadImage]
+    () => ({
+      $progressEvent,
+      $progressImage,
+      $hasProgressImage,
+      $isProgressImageResolving,
+      $isTemporarilyShowingSelectedImage,
+      onLoadImage,
+    }),
+    [
+      $hasProgressImage,
+      $isProgressImageResolving,
+      $isTemporarilyShowingSelectedImage,
+      $progressEvent,
+      $progressImage,
+      onLoadImage,
+    ]
   );
 
   return <ImageViewerContext.Provider value={value}>{props.children}</ImageViewerContext.Provider>;