diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentImagePreview.tsx b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentImagePreview.tsx index 1978a7fc1ab..ba76ca299bb 100644 --- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentImagePreview.tsx +++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentImagePreview.tsx @@ -193,6 +193,12 @@ export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | nu dependencies: [onHotkeyNextImage], }); + // The loaded image identifies its session so the viewer can tell a late load from an earlier + // session apart from the one whose preview is currently retained (see onLoadImage). + const onLoadRenderedImage = useCallback(() => { + onLoadImage(imageToRender?.session_id ?? null); + }, [imageToRender?.session_id, onLoadImage]); + const withProgress = shouldShowProgressInViewer && hasProgressImage && !isTemporarilyShowingSelectedImage; // When more than one session is generating concurrently (multi-GPU), tile their previews instead of // showing only the most recent one. @@ -210,7 +216,7 @@ export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | nu > {imageToRender && ( - + )} {!imageToRender && } diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentVideoPreview.tsx b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentVideoPreview.tsx index 9329a7173c7..19aacb1944b 100644 --- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentVideoPreview.tsx +++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentVideoPreview.tsx @@ -302,7 +302,7 @@ export const CurrentVideoPreview = memo(({ videoDTO }: Props) => { // video frame until playback or a seek — the element just shows its black background. // Setting currentTime to 0.0001 nudges the decoder to paint without measurably advancing. const handleLoadedMetadata = useCallback(() => { - onLoadImage(); + onLoadImage(videoDTO?.session_id ?? null); const el = videoRef.current; if (el && !isPlaying && el.currentTime === 0) { try { @@ -311,7 +311,7 @@ export const CurrentVideoPreview = memo(({ videoDTO }: Props) => { // Some browsers throw if metadata isn't fully ready yet; harmless. } } - }, [isPlaying, onLoadImage]); + }, [isPlaying, onLoadImage, videoDTO?.session_id]); if (!videoDTO) { return ; diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/context.tsx b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/context.tsx index 74fb418761f..19530ffebb9 100644 --- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/context.tsx +++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/context.tsx @@ -1,27 +1,25 @@ import { useStore } from '@nanostores/react'; import { logger } from 'app/logging/logger'; import { useAppSelector, useAppStore } from 'app/store/storeHooks'; +import { selectCurrentUser } from 'features/auth/store/authSlice'; +import type { + ViewerProgressDataMap, + ViewerProgressDatum, +} from 'features/gallery/components/ImageViewer/viewerProgressLifecycle'; +import { createViewerProgressLifecycle } from 'features/gallery/components/ImageViewer/viewerProgressLifecycle'; 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, map, type MapStore, type WritableAtom } from 'nanostores'; import type { PropsWithChildren } from 'react'; -import { createContext, memo, useCallback, useContext, useEffect, useMemo, useRef, useState } from 'react'; +import { createContext, memo, useCallback, useContext, useEffect, useMemo, useState } from 'react'; import type { S } from 'services/api/types'; import { getEventScope } from 'services/events/eventScope'; import { $socket } from 'services/events/stores'; import { assert } from 'tsafe'; import type { JsonObject } from 'type-fest'; -/** Live progress for a single in-flight session (queue item). Used to tile the viewer when several - * sessions run concurrently (multi-GPU). Only items that have produced a preview image are tracked. */ -export type ViewerProgressDatum = { - itemId: number; - progressEvent: S['InvocationProgressEvent']; - progressImage: ProgressImageType; -}; - -type ViewerProgressDataMap = Record; +export type { ViewerProgressDatum } from 'features/gallery/components/ImageViewer/viewerProgressLifecycle'; type ImageViewerContextValue = { $progressEvent: Atom; @@ -33,7 +31,11 @@ type ImageViewerContextValue = { $activeProgressData: Atom; $isProgressImageResolving: Atom; $isTemporarilyShowingSelectedImage: WritableAtom; - onLoadImage: () => void; + /** + * The viewer finished loading the final image/video for the given session (its DTO's + * `session_id`, or null when it has none). Ends the completed session's "resolve" illusion. + */ + onLoadImage: (sessionId: string | null) => void; }; const ImageViewerContext = createContext(null); @@ -58,10 +60,24 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => { )[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({ max: 200 })); + // Session id -> queue item id, learned from progress events. Outlives the item's terminal event + // so a late final-image load can be attributed to the session that produced it. + const [itemIdBySessionId] = useState(() => new LRUCache({ max: 200 })); + // All store mutations live in the lifecycle (extracted for unit testing); the effects below own + // the socket subscriptions and the ownership/scope checks on incoming events. + const lifecycle = useState(() => + createViewerProgressLifecycle({ + $progressEvent, + $progressImage, + $progressData, + $isProgressImageResolving, + finishedQueueItemIds, + itemIdBySessionId, + }) + )[0]; useEffect(() => { if (!socket) { @@ -74,24 +90,11 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => { if (getEventScope(store.getState, data) !== 'own') { return; } - if (finishedQueueItemIds.has(data.item_id)) { + if (!lifecycle.recordProgress(data)) { log.trace( { data } as JsonObject, `Received InvocationProgressEvent event for already-finished queue item ${data.item_id}` ); - return; - } - shouldClearProgressImageOnLoadRef.current = false; - $isProgressImageResolving.set(false); - $progressEvent.set(data); - if (data.image) { - $progressImage.set(data.image); - // Track per-session so the viewer can tile concurrent sessions (multi-GPU). - $progressData.setKey(data.item_id, { - itemId: data.item_id, - progressEvent: data, - progressImage: data.image, - }); } }; @@ -100,7 +103,7 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => { return () => { socket.off('invocation_progress', onInvocationProgress); }; - }, [$isProgressImageResolving, $progressData, $progressEvent, $progressImage, finishedQueueItemIds, socket, store]); + }, [lifecycle, socket, store]); useEffect(() => { if (!socket) { @@ -116,51 +119,22 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => { if (getEventScope(store.getState, data) !== 'own') { return; } - if (finishedQueueItemIds.has(data.item_id)) { + if (data.status === 'in_progress') { + // Track the claim itself, not just the progress events that follow it: a queue clear + // cancels the items already running before it deletes the rows, so a worker that claims an + // item in between never gets a terminal event and its first progress event lands after the + // clear (see the lifecycle's onItemStarted). + lifecycle.onItemStarted(data.item_id); + return; + } + if (data.status !== 'completed' && data.status !== 'canceled' && data.status !== 'failed') { + return; + } + if (!lifecycle.onTerminal(data, autoSwitch)) { log.trace( { data } as JsonObject, `Received QueueItemStatusChangedEvent event for already-finished queue item ${data.item_id}` ); - return; - } - if (data.status === 'completed' || data.status === 'canceled' || data.status === 'failed') { - finishedQueueItemIds.set(data.item_id, true); - // Remove this session's tile from the multi-session preview as soon as it reaches a terminal - // state. The single-image "resolve" illusion below is handled separately via onLoadImage. - $progressData.setKey(data.item_id, undefined); - // The shared $progressEvent/$progressImage globals may currently hold a DIFFERENT session's - // latest preview (multi-GPU). Only the item that owns them may clear them — otherwise - // canceling item A would blank item B's still-running preview until B's next image event. - const globalProgressEvent = $progressEvent.get(); - if (globalProgressEvent !== null && globalProgressEvent.item_id !== data.item_id) { - return; - } - // Completed queue items have the progress event cleared by the onLoadImage callback. This allows the viewer to - // create the illusion of the progress image "resolving" into the final image. If we cleared the progress image - // now, there would be a flicker where the progress image disappears before the final image appears, and the - // last-selected gallery image should be shown for a brief moment. - // - // When gallery auto-switch is disabled, we do not need to create this illusion, because we are not going to - // switch to the final image automatically. In this case, we clear the progress image immediately. - // - // We also clear the progress image if the queue item is canceled or failed, as there is no final image to show. - if ( - data.status === 'canceled' || - data.status === 'failed' || - !autoSwitch || - // When the origin is 'canvas' and destination is 'canvas' (without a ':' suffix), that means the - // image is going to be added to the staging area. In this case, we need to clear the progress image else it - // 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); - } } }; @@ -169,27 +143,55 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => { return () => { socket.off('queue_item_status_changed', onQueueItemStatusChanged); }; - }, [ - $isProgressImageResolving, - $progressData, - $progressEvent, - $progressImage, - autoSwitch, - finishedQueueItemIds, - socket, - store, - ]); - - const onLoadImage = useCallback(() => { - if (!shouldClearProgressImageOnLoadRef.current) { + }, [autoSwitch, lifecycle, socket, store]); + + useEffect(() => { + if (!socket) { + return; + } + + const onQueueCleared = (data: S['QueueClearedEvent']) => { + // Scope is decided inside the lifecycle: it needs the current user id to tell whether the + // clear could have deleted this client's items (see onQueueCleared's docstring). + const currentUserId = selectCurrentUser(store.getState())?.user_id ?? null; + lifecycle.onQueueCleared(data, currentUserId); + }; + + socket.on('queue_cleared', onQueueCleared); + + return () => { + socket.off('queue_cleared', onQueueCleared); + }; + }, [lifecycle, socket, store]); + + useEffect(() => { + if (!socket) { return; } - shouldClearProgressImageOnLoadRef.current = false; - $isProgressImageResolving.set(false); - $progressEvent.set(null); - $progressImage.set(null); - }, [$isProgressImageResolving, $progressEvent, $progressImage]); + const onDisconnect = () => { + // Mirrors the app-wide progress stores (see setEventListeners): a disconnected socket may + // miss terminal events, so previews from before the gap cannot be trusted to ever resolve. + lifecycle.reset(); + }; + + socket.on('disconnect', onDisconnect); + + return () => { + socket.off('disconnect', onDisconnect); + // The socket is being replaced (e.g. an auth-token change swapped $socket for a different + // user's connection) or the viewer is unmounting: any tracked sessions belong to the old + // connection and will never emit another terminal event here. + lifecycle.reset(); + }; + }, [lifecycle, socket]); + + const onLoadImage = useCallback( + (sessionId: string | null) => { + lifecycle.onFinalImageLoaded(sessionId); + }, + [lifecycle] + ); const value = useMemo( () => ({ diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.test.ts b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.test.ts new file mode 100644 index 00000000000..a3181382a7f --- /dev/null +++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.test.ts @@ -0,0 +1,357 @@ +import type { ProgressImage as ProgressImageType } from 'features/nodes/types/common'; +import { atom, map } from 'nanostores'; +import type { S } from 'services/api/types'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +import type { ViewerProgressDataMap, ViewerProgressStores } from './viewerProgressLifecycle'; +import { createViewerProgressLifecycle, RESOLVE_TIMEOUT_MS } from './viewerProgressLifecycle'; + +const buildProgressImage = (itemId: number): ProgressImageType => + ({ + dataURL: `data:image/png;base64,item-${itemId}`, + width: 512, + height: 512, + }) as ProgressImageType; + +const buildProgressEvent = (overrides: Partial = {}): S['InvocationProgressEvent'] => + ({ + queue_id: 'default', + item_id: 1, + batch_id: 'batch-1', + origin: null, + destination: null, + user_id: 'user-1', + session_id: 'session-1', + invocation_source_id: 'node-1', + invocation: { id: 'node-1', type: 'test_node' }, + message: 'denoising', + percentage: 0.5, + image: null, + ...overrides, + }) as S['InvocationProgressEvent']; + +const buildTerminalEvent = ( + overrides: Partial = {} +): S['QueueItemStatusChangedEvent'] => + ({ + queue_id: 'default', + item_id: 1, + batch_id: 'batch-1', + origin: null, + destination: null, + user_id: 'user-1', + status: 'completed', + ...overrides, + }) as S['QueueItemStatusChangedEvent']; + +const buildQueueClearedEvent = (userId: string | null): S['QueueClearedEvent'] => + ({ queue_id: 'default', user_id: userId }) as S['QueueClearedEvent']; + +describe('viewerProgressLifecycle', () => { + let stores: ViewerProgressStores; + let lifecycle: ReturnType; + + beforeEach(() => { + vi.useFakeTimers(); + stores = { + $progressEvent: atom(null), + $progressImage: atom(null), + $progressData: map({}), + $isProgressImageResolving: atom(false), + finishedQueueItemIds: new Map(), + itemIdBySessionId: new Map(), + }; + lifecycle = createViewerProgressLifecycle(stores); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + const startTwoSessions = () => { + // B posts a preview first, then A — A owns the shared single-image preview. + const eventB = buildProgressEvent({ item_id: 2, session_id: 'session-2', image: buildProgressImage(2) }); + const eventA = buildProgressEvent({ item_id: 1, session_id: 'session-1', image: buildProgressImage(1) }); + lifecycle.recordProgress(eventB); + lifecycle.recordProgress(eventA); + return { eventA, eventB }; + }; + + describe('recordProgress', () => { + it('tracks per-session data and hands the shared preview to the latest reporter', () => { + const { eventA } = startTwoSessions(); + expect(stores.$progressEvent.get()).toBe(eventA); + expect(stores.$progressImage.get()).toBe(eventA.image); + expect(stores.$progressData.get()[1]?.itemId).toBe(1); + expect(stores.$progressData.get()[2]?.itemId).toBe(2); + }); + + it('ignores events for finished items so trailing progress cannot repopulate the preview', () => { + stores.finishedQueueItemIds.set(1, true); + const handled = lifecycle.recordProgress(buildProgressEvent({ item_id: 1, image: buildProgressImage(1) })); + expect(handled).toBe(false); + expect(stores.$progressEvent.get()).toBeNull(); + expect(stores.$progressData.get()[1]).toBeUndefined(); + }); + }); + + describe('onTerminal', () => { + it.each(['completed', 'canceled', 'failed'] as const)( + 'hands the shared preview to the remaining session when its owner reaches %s', + (status) => { + const { eventB } = startTwoSessions(); + lifecycle.onTerminal(buildTerminalEvent({ item_id: 1, status }), true); + // B immediately becomes the single visible progress image and indicator. + expect(stores.$progressEvent.get()).toBe(eventB); + expect(stores.$progressImage.get()).toBe(eventB.image); + expect(stores.$progressData.get()[1]).toBeUndefined(); + expect(stores.$progressData.get()[2]?.itemId).toBe(2); + // No resolve illusion may be pending — it would swap B's live preview for A's final image. + expect(stores.$isProgressImageResolving.get()).toBe(false); + lifecycle.onFinalImageLoaded('session-1'); + expect(stores.$progressEvent.get()).toBe(eventB); + expect(stores.$progressImage.get()).toBe(eventB.image); + } + ); + + it('hands the shared preview to the remaining session even when auto-switch is off', () => { + const { eventB } = startTwoSessions(); + lifecycle.onTerminal(buildTerminalEvent({ item_id: 1, status: 'canceled' }), false); + expect(stores.$progressEvent.get()).toBe(eventB); + expect(stores.$progressImage.get()).toBe(eventB.image); + }); + + it('promotes the most recently updated remaining session when several remain', () => { + startTwoSessions(); + const eventC = buildProgressEvent({ item_id: 3, session_id: 'session-3', image: buildProgressImage(3) }); + lifecycle.recordProgress(eventC); + lifecycle.onTerminal(buildTerminalEvent({ item_id: 3, status: 'canceled' }), true); + // C owned the preview; A reported more recently than B, so A takes over. + expect(stores.$progressEvent.get()?.item_id).toBe(1); + }); + + it('leaves the shared preview alone when a non-owner terminates', () => { + const { eventA } = startTwoSessions(); + lifecycle.onTerminal(buildTerminalEvent({ item_id: 2, status: 'canceled' }), true); + expect(stores.$progressEvent.get()).toBe(eventA); + expect(stores.$progressImage.get()).toBe(eventA.image); + expect(stores.$progressData.get()[2]).toBeUndefined(); + }); + + it('clears immediately when the last session is canceled', () => { + const eventA = buildProgressEvent({ item_id: 1, image: buildProgressImage(1) }); + lifecycle.recordProgress(eventA); + lifecycle.onTerminal(buildTerminalEvent({ item_id: 1, status: 'canceled' }), true); + expect(stores.$progressEvent.get()).toBeNull(); + expect(stores.$progressImage.get()).toBeNull(); + expect(stores.$isProgressImageResolving.get()).toBe(false); + }); + + it('runs the resolve illusion when the last session completes with auto-switch on', () => { + const eventA = buildProgressEvent({ item_id: 1, image: buildProgressImage(1) }); + lifecycle.recordProgress(eventA); + lifecycle.onTerminal(buildTerminalEvent({ item_id: 1, status: 'completed' }), true); + // The preview is retained until the final image loads, "resolving" into it. + expect(stores.$progressEvent.get()).toBe(eventA); + expect(stores.$isProgressImageResolving.get()).toBe(true); + lifecycle.onFinalImageLoaded('session-1'); + expect(stores.$progressEvent.get()).toBeNull(); + expect(stores.$progressImage.get()).toBeNull(); + expect(stores.$isProgressImageResolving.get()).toBe(false); + }); + + it('runs no resolve illusion when the completed item never reported progress', () => { + // Nothing is retained, so arming the illusion would only leave the resolving flag stuck on. + lifecycle.onTerminal(buildTerminalEvent({ item_id: 1, status: 'completed' }), true); + expect(stores.$isProgressImageResolving.get()).toBe(false); + expect(stores.$progressEvent.get()).toBeNull(); + }); + + it('ignores repeat terminal events for an already-finished item', () => { + const { eventA } = startTwoSessions(); + lifecycle.onTerminal(buildTerminalEvent({ item_id: 2, status: 'canceled' }), true); + expect(lifecycle.onTerminal(buildTerminalEvent({ item_id: 2, status: 'canceled' }), true)).toBe(false); + expect(stores.$progressEvent.get()).toBe(eventA); + }); + }); + + describe('onFinalImageLoaded', () => { + it('ignores a late load from a session that finished before the current preview owner', () => { + const { eventB } = startTwoSessions(); + // A completes first and hands the shared preview to the still-running B... + lifecycle.onTerminal(buildTerminalEvent({ item_id: 1, status: 'completed' }), true); + // ...then B completes and starts its own resolve illusion, retaining B's last frame. + lifecycle.onTerminal(buildTerminalEvent({ item_id: 2, status: 'completed' }), true); + expect(stores.$isProgressImageResolving.get()).toBe(true); + // A's final image only now finishes loading. Clearing here would cut B's illusion short. + lifecycle.onFinalImageLoaded('session-1'); + expect(stores.$progressEvent.get()).toBe(eventB); + expect(stores.$progressImage.get()).toBe(eventB.image); + expect(stores.$isProgressImageResolving.get()).toBe(true); + // B's own final image ends the illusion. + lifecycle.onFinalImageLoaded('session-2'); + expect(stores.$progressEvent.get()).toBeNull(); + expect(stores.$progressImage.get()).toBeNull(); + expect(stores.$isProgressImageResolving.get()).toBe(false); + }); + + it('clears the retained preview on a timeout when its final image never loads', () => { + // The retained session's image may never load: the load can fail (the viewer's + // reports errors through onError, not onLoad), or auto-switch may end up selecting a + // concurrently-completed session's image, so the retained session's image is never + // rendered. An ignored load may have been the last one coming, so the preview — an opaque + // overlay — must not be left covering the viewer until the next generation. + const { eventB } = startTwoSessions(); + lifecycle.onTerminal(buildTerminalEvent({ item_id: 1, status: 'completed' }), true); + lifecycle.onTerminal(buildTerminalEvent({ item_id: 2, status: 'completed' }), true); + lifecycle.onFinalImageLoaded('session-1'); + expect(stores.$progressEvent.get()).toBe(eventB); + vi.advanceTimersByTime(RESOLVE_TIMEOUT_MS); + expect(stores.$progressEvent.get()).toBeNull(); + expect(stores.$progressImage.get()).toBeNull(); + expect(stores.$isProgressImageResolving.get()).toBe(false); + }); + + it.each([ + [ + 'a still-running session takes the preview over', + () => { + lifecycle.recordProgress( + buildProgressEvent({ item_id: 2, session_id: 'session-2', image: buildProgressImage(2) }) + ); + }, + ], + [ + 'the final image loads first', + () => { + lifecycle.onFinalImageLoaded('session-1'); + }, + ], + [ + 'the preview state is reset', + () => { + lifecycle.reset(); + }, + ], + ])('disarms the resolve timeout when %s', (_desc, takeOver) => { + const eventA = buildProgressEvent({ item_id: 1, image: buildProgressImage(1) }); + lifecycle.recordProgress(eventA); + lifecycle.onTerminal(buildTerminalEvent({ item_id: 1, status: 'completed' }), true); + expect(vi.getTimerCount()).toBe(1); + takeOver(); + // The timeout must never fire against a preview that has since been replaced or dropped. + // Asserting the timer is gone, not just that the state survives it, is what makes this bite + // for the takeovers that leave the stores null — there, a leaked timer would clear state + // that is already clear. + expect(vi.getTimerCount()).toBe(0); + const eventAfterTakeOver = stores.$progressEvent.get(); + vi.advanceTimersByTime(RESOLVE_TIMEOUT_MS * 2); + expect(stores.$progressEvent.get()).toBe(eventAfterTakeOver); + }); + + it.each([ + ['an image with no session (e.g. an upload)', null], + ['an image from a session this viewer never tracked', 'session-from-a-previous-visit'], + ])('still clears the retained preview on %s', (_desc, sessionId) => { + const eventA = buildProgressEvent({ item_id: 1, image: buildProgressImage(1) }); + lifecycle.recordProgress(eventA); + lifecycle.onTerminal(buildTerminalEvent({ item_id: 1, status: 'completed' }), true); + expect(stores.$isProgressImageResolving.get()).toBe(true); + // Unattributable loads keep the safety net: a retained preview must not cover the viewer + // indefinitely just because the completed item's own image never loads. + lifecycle.onFinalImageLoaded(sessionId); + expect(stores.$progressEvent.get()).toBeNull(); + expect(stores.$progressImage.get()).toBeNull(); + expect(stores.$isProgressImageResolving.get()).toBe(false); + }); + + it('clears on a load from a session tracked before a reset', () => { + // Attributions are dropped along with the state they describe, so images generated before a + // disconnect or socket swap can still end a later session's illusion. + startTwoSessions(); + lifecycle.reset(); + const eventC = buildProgressEvent({ item_id: 3, session_id: 'session-3', image: buildProgressImage(3) }); + lifecycle.recordProgress(eventC); + lifecycle.onTerminal(buildTerminalEvent({ item_id: 3, status: 'completed' }), true); + expect(stores.$isProgressImageResolving.get()).toBe(true); + lifecycle.onFinalImageLoaded('session-1'); + expect(stores.$progressEvent.get()).toBeNull(); + expect(stores.$isProgressImageResolving.get()).toBe(false); + }); + }); + + describe('onQueueCleared', () => { + it.each([ + ['an unscoped (admin or single-user) clear', null, 'user-1'], + ["the current user's scoped clear", 'user-1', 'user-1'], + ])('drops all previews and blocks trailing progress on %s', (_desc, clearedUserId, currentUserId) => { + startTwoSessions(); + const applied = lifecycle.onQueueCleared(buildQueueClearedEvent(clearedUserId), currentUserId); + expect(applied).toBe(true); + expect(stores.$progressEvent.get()).toBeNull(); + expect(stores.$progressImage.get()).toBeNull(); + expect(stores.$progressData.get()).toEqual({}); + // A worker claimed between the clear's cancellation pass and deletion is stopped only by + // this event — its trailing progress must not repopulate the preview. + expect(lifecycle.recordProgress(buildProgressEvent({ item_id: 1, image: buildProgressImage(1) }))).toBe(false); + expect(lifecycle.recordProgress(buildProgressEvent({ item_id: 2, image: buildProgressImage(2) }))).toBe(false); + }); + + it('blocks trailing progress from an item claimed after the clear cancelled the running ones', () => { + // The clear cancels the items already running before deleting the rows, so a worker that + // claims an item in between never gets a terminal event: only its in_progress claim, which + // precedes the deletion, tells the viewer the item exists. Its first progress event arrives + // seconds later — a preview for a deleted item that nothing would ever take down. + expect(lifecycle.onItemStarted(7)).toBe(true); + expect(lifecycle.onQueueCleared(buildQueueClearedEvent(null), 'user-1')).toBe(true); + expect( + lifecycle.recordProgress( + buildProgressEvent({ item_id: 7, session_id: 'session-7', image: buildProgressImage(7) }) + ) + ).toBe(false); + expect(stores.$progressEvent.get()).toBeNull(); + expect(stores.$progressImage.get()).toBeNull(); + }); + + it('blocks trailing progress from a session that had not produced a preview image yet', () => { + // Item 1 has only reported progress without an image, so it is absent from $progressData and + // does not own the shared globals once item 2 reports an image. + lifecycle.recordProgress(buildProgressEvent({ item_id: 1, session_id: 'session-1', image: null })); + lifecycle.recordProgress( + buildProgressEvent({ item_id: 2, session_id: 'session-2', image: buildProgressImage(2) }) + ); + expect(lifecycle.onQueueCleared(buildQueueClearedEvent(null), 'user-1')).toBe(true); + // Its first image event must not resurrect the preview the clear just dropped. + expect(lifecycle.recordProgress(buildProgressEvent({ item_id: 1, image: buildProgressImage(1) }))).toBe(false); + expect(stores.$progressEvent.get()).toBeNull(); + expect(stores.$progressImage.get()).toBeNull(); + expect(stores.$progressData.get()).toEqual({}); + }); + + it.each([ + ["another user's scoped clear", 'user-2'], + ['the sanitized broadcast of a foreign scoped clear', 'redacted'], + ])('leaves previews alone on %s', (_desc, clearedUserId) => { + const { eventA } = startTwoSessions(); + const applied = lifecycle.onQueueCleared(buildQueueClearedEvent(clearedUserId), 'user-1'); + expect(applied).toBe(false); + expect(stores.$progressEvent.get()).toBe(eventA); + expect(stores.$progressData.get()[1]?.itemId).toBe(1); + expect(stores.$progressData.get()[2]?.itemId).toBe(2); + }); + }); + + describe('reset', () => { + it('drops all preview state without marking items finished', () => { + startTwoSessions(); + stores.$isProgressImageResolving.set(true); + lifecycle.reset(); + expect(stores.$progressEvent.get()).toBeNull(); + expect(stores.$progressImage.get()).toBeNull(); + expect(stores.$progressData.get()).toEqual({}); + expect(stores.$isProgressImageResolving.get()).toBe(false); + // A new connection's events for a re-used id are not blocked — reset is not a cancel. + expect(lifecycle.recordProgress(buildProgressEvent({ item_id: 1, image: buildProgressImage(1) }))).toBe(true); + }); + }); +}); diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.ts b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.ts new file mode 100644 index 00000000000..e6c78cc7499 --- /dev/null +++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.ts @@ -0,0 +1,319 @@ +import type { ProgressImage as ProgressImageType } from 'features/nodes/types/common'; +import type { MapStore, WritableAtom } from 'nanostores'; +import type { S } from 'services/api/types'; + +/** Live progress for a single in-flight session (queue item). Used to tile the viewer when several + * sessions run concurrently (multi-GPU). Only items that have produced a preview image are tracked. + * `seq` orders data by most recent update, so the shared single-image preview can be handed to the + * freshest remaining session when its current owner terminates. */ +export type ViewerProgressDatum = { + itemId: number; + seq: number; + progressEvent: S['InvocationProgressEvent']; + progressImage: ProgressImageType; +}; + +export type ViewerProgressDataMap = Record; + +/** The subset of LRUCache the lifecycle needs — kept minimal so tests can pass a plain Map. */ +type FinishedQueueItemIds = { + has: (itemId: number) => boolean; + set: (itemId: number, value: boolean) => unknown; +}; + +/** The subset of LRUCache the lifecycle needs — kept minimal so tests can pass a plain Map. */ +type ItemIdBySessionId = { + get: (sessionId: string) => number | undefined; + set: (sessionId: string, itemId: number) => unknown; + clear: () => unknown; +}; + +/** + * How long a completed session's retained preview may wait for its final image to load before it + * is cleared anyway. The "resolve" illusion normally ends on that image's load event, but nothing + * guarantees the event arrives: the load can fail (the viewer's reports errors through + * onError, not onLoad), or auto-switch can end up selecting a concurrently-completed session's + * image instead, in which case the retained session's image is never rendered at all. Without this + * bound the preview — an opaque overlay — could cover the viewer until the next generation. + */ +export const RESOLVE_TIMEOUT_MS = 3000; + +export type ViewerProgressStores = { + $progressEvent: WritableAtom; + $progressImage: WritableAtom; + /** Per-session progress, keyed by queue item id. Drives the tiled multi-session preview. */ + $progressData: MapStore; + $isProgressImageResolving: WritableAtom; + /** Finished queue items, tracked so trailing progress events cannot repopulate the preview. */ + finishedQueueItemIds: FinishedQueueItemIds; + /** Queue item id of each session we have seen progress for, keyed by session id. Outlives the + * item's terminal event so a late final-image load can be attributed to the session that + * produced it (see onFinalImageLoaded). */ + itemIdBySessionId: ItemIdBySessionId; +}; + +const pickLatestDatum = (data: ViewerProgressDataMap): ViewerProgressDatum | null => { + let latest: ViewerProgressDatum | null = null; + for (const datum of Object.values(data)) { + if (datum !== undefined && (latest === null || datum.seq > latest.seq)) { + latest = datum; + } + } + return latest; +}; + +/** + * The store-side lifecycle of the image viewer's live-preview state, factored out of the React + * provider so it can be unit tested. The provider owns the socket subscriptions and the + * ownership/scope checks on incoming events; every store mutation happens here. + * + * The state it manages: + * - `$progressData`: one entry per session with a preview image (the tiled multi-session view). + * - `$progressEvent` / `$progressImage`: the shared single-image preview, owned by the session + * that most recently reported progress. + */ +export const createViewerProgressLifecycle = (stores: ViewerProgressStores) => { + const { + $progressEvent, + $progressImage, + $progressData, + $isProgressImageResolving, + finishedQueueItemIds, + itemIdBySessionId, + } = stores; + let seq = 0; + // The queue item whose retained preview the final gallery image's onLoad should clear — the tail + // end of the "resolve" illusion for a completed session (see onTerminal / onFinalImageLoaded). + // Null when no illusion is pending. Always written through setPendingResolve, which keeps the + // safety timeout in sync. + let pendingResolveItemId: number | null = null; + // Every item we have seen progress for and not yet seen terminate, including items that have not + // produced a preview image (those are absent from $progressData). A queue clear deletes items + // without emitting a per-item terminal event, so this is the set that must be marked finished + // there — otherwise an image-less session could later emit an image and resurrect the preview. + const unfinishedItemIds = new Set(); + let resolveTimeoutId: ReturnType | null = null; + + const clearRetainedPreview = (): void => { + $isProgressImageResolving.set(false); + $progressEvent.set(null); + $progressImage.set(null); + }; + + /** + * Arm (or, with null, disarm) the pending "resolve" illusion. Every write to + * `pendingResolveItemId` goes through here, so the safety timeout exists exactly while the + * illusion is pending: anything that ends the illusion — a load, a takeover by another session, + * a reset — also cancels the timeout, and it can never clear a preview that some other session + * has since taken over. + * + * The timeout bounds the illusion only. A preview left standing by a path that does not arm one + * (a progress event that carries no image replaces $progressEvent but not $progressImage, so the + * previous session's frame stays up while the next queue item spins up) is unbounded here, as it + * is today — that frame is taken down by the next preview image rather than by this timeout. + */ + const setPendingResolve = (itemId: number | null): void => { + if (resolveTimeoutId !== null) { + clearTimeout(resolveTimeoutId); + resolveTimeoutId = null; + } + pendingResolveItemId = itemId; + if (itemId === null) { + return; + } + resolveTimeoutId = setTimeout(() => { + resolveTimeoutId = null; + pendingResolveItemId = null; + clearRetainedPreview(); + }, RESOLVE_TIMEOUT_MS); + }; + + const clearAll = (): void => { + setPendingResolve(null); + unfinishedItemIds.clear(); + // Session attributions describe state this reset just dropped. Keeping them would make later + // loads of those images look like another session's, suppressing clears they should perform. + itemIdBySessionId.clear(); + clearRetainedPreview(); + $progressData.set({}); + }; + + /** + * A worker claimed this queue item (`in_progress`). Tracked so a queue clear can mark it + * finished: the clear cancels the items that were already running before it deletes the rows, + * but a worker that claims an item in between gets no terminal event at all — its row is gone — + * and its first progress event would otherwise appear seconds after the clear and put a preview + * for a deleted item on screen, with nothing left to ever take it down. + * + * Returns false if the item already finished (event ignored). + */ + const onItemStarted = (itemId: number): boolean => { + if (finishedQueueItemIds.has(itemId)) { + return false; + } + unfinishedItemIds.add(itemId); + return true; + }; + + /** Record a progress event. Returns false if the item already finished (event ignored). */ + const recordProgress = (data: S['InvocationProgressEvent']): boolean => { + if (finishedQueueItemIds.has(data.item_id)) { + return false; + } + unfinishedItemIds.add(data.item_id); + itemIdBySessionId.set(data.session_id, data.item_id); + setPendingResolve(null); + $isProgressImageResolving.set(false); + $progressEvent.set(data); + if (data.image) { + $progressImage.set(data.image); + // Track per-session so the viewer can tile concurrent sessions (multi-GPU). + $progressData.setKey(data.item_id, { + itemId: data.item_id, + seq: ++seq, + progressEvent: data, + progressImage: data.image, + }); + } + return true; + }; + + /** Handle a terminal status for a queue item. Returns false if it already finished (ignored). */ + const onTerminal = (data: S['QueueItemStatusChangedEvent'], autoSwitch: boolean): boolean => { + if (finishedQueueItemIds.has(data.item_id)) { + return false; + } + finishedQueueItemIds.set(data.item_id, true); + unfinishedItemIds.delete(data.item_id); + // Remove this session's tile from the multi-session preview as soon as it reaches a terminal + // state. The single-image "resolve" illusion below is handled separately via onLoadImage. + $progressData.setKey(data.item_id, undefined); + // The shared $progressEvent/$progressImage globals may currently hold a DIFFERENT session's + // latest preview (multi-GPU). Only the item that owns them may replace or clear them — + // otherwise canceling item A would blank item B's still-running preview until B's next image + // event. + const globalProgressEvent = $progressEvent.get(); + if (globalProgressEvent !== null && globalProgressEvent.item_id !== data.item_id) { + return true; + } + const successor = pickLatestDatum($progressData.get()); + if (successor !== null) { + // The terminated item owned the shared preview, but other sessions are still generating: + // hand the preview to the most recently updated one immediately. The tiled view only renders + // with more than one active session, so once a single session remains it is displayed + // through these globals — leaving them cleared (or parked on the finished session's stale + // frame via the resolve illusion) would hide a still-running preview. This applies to every + // terminal status, including successful completion with auto-switch. + setPendingResolve(null); + $isProgressImageResolving.set(false); + $progressEvent.set(successor.progressEvent); + $progressImage.set(successor.progressImage); + return true; + } + if (globalProgressEvent === null) { + // Nothing is retained (this item never reported progress), so there is no illusion to run — + // arming one would leave $isProgressImageResolving stuck on until the next generation. + setPendingResolve(null); + $isProgressImageResolving.set(false); + return true; + } + // Completed queue items have the progress event cleared by the onLoadImage callback. This allows the viewer to + // create the illusion of the progress image "resolving" into the final image. If we cleared the progress image + // now, there would be a flicker where the progress image disappears before the final image appears, and the + // last-selected gallery image should be shown for a brief moment. + // + // When gallery auto-switch is disabled, we do not need to create this illusion, because we are not going to + // switch to the final image automatically. In this case, we clear the progress image immediately. + // + // We also clear the progress image if the queue item is canceled or failed, as there is no final image to show. + if ( + data.status === 'canceled' || + data.status === 'failed' || + !autoSwitch || + // When the origin is 'canvas' and destination is 'canvas' (without a ':' suffix), that means the + // image is going to be added to the staging area. In this case, we need to clear the progress image else it + // will be stuck on the viewer. + (data.origin === 'canvas' && data.destination !== 'canvas') + ) { + setPendingResolve(null); + clearRetainedPreview(); + } else { + setPendingResolve(data.item_id); + $isProgressImageResolving.set(true); + } + return true; + }; + + /** + * The final gallery image (or video) finished loading. If a completed session's "resolve" + * illusion is pending, this is its tail end: the retained preview is cleared so the final image + * shows. A no-op otherwise (e.g. when the preview was handed to a still-running session). + * + * `sessionId` identifies the item that was loaded (ImageDTO/VideoDTO `session_id`). Several + * sessions run concurrently under multi-GPU and auto-switch, so a load can arrive late, after a + * *different* session took over the retained preview: session A completes and hands the preview + * to B, B then completes and starts its own resolve illusion, and only then does A's final image + * finish loading. Clearing on A's load would cut B's illusion short — exactly the flicker the + * illusion exists to hide — so a load is ignored when it can be positively attributed to another + * session we tracked. Loads we cannot attribute (uploads, images from before this viewer + * mounted) still clear. + * + * Ignoring a load must never be the difference between the preview clearing and not clearing: + * the retained session's own image may never load at all (see RESOLVE_TIMEOUT_MS), and the + * ignored load may have been the last one coming. The timeout armed alongside the illusion is + * what bounds it — this check only decides whether the illusion ends early. + */ + const onFinalImageLoaded = (sessionId: string | null): void => { + if (pendingResolveItemId === null) { + return; + } + if (sessionId !== null) { + const loadedItemId = itemIdBySessionId.get(sessionId); + if (loadedItemId !== undefined && loadedItemId !== pendingResolveItemId) { + return; + } + } + setPendingResolve(null); + clearRetainedPreview(); + }; + + /** + * Handle a queue-cleared event. A clear deletes queue items without emitting a per-item terminal + * status event for every one of them (a worker claimed mid-clear is stopped only by this event), + * so the tracked previews must be dropped here. Which items were deleted depends on the event's + * scope (mirroring workflowExecutionCoordinator.onQueueCleared): an unscoped clear (user_id=null + * — an admin or single-user clear) deleted every item; a clear scoped to the current user + * deleted all of this client's items; another user's scoped clear — received in full by admins + * or as the sanitized user_id="redacted" broadcast by everyone else — deleted none of this + * client's items, and this store only ever tracks the client's own items. + * + * Returns whether the clear applied to this client's previews. + */ + const onQueueCleared = (data: S['QueueClearedEvent'], currentUserId: string | null): boolean => { + const clearedUserId = data.user_id ?? null; + if (clearedUserId !== null && clearedUserId !== currentUserId) { + return false; + } + // Mark every session we have seen progress for and not yet seen terminate as finished, so a + // trailing invocation_progress event from a worker that the clear is still stopping cannot + // repopulate the preview. This must cover sessions that have not produced a preview image yet + // — they are absent from $progressData and may not own the shared globals, but their first + // image event would otherwise resurrect the preview after the clear. + for (const itemId of unfinishedItemIds) { + finishedQueueItemIds.set(itemId, true); + } + clearAll(); + return true; + }; + + /** + * Drop all preview state without marking items finished. For socket disconnection and socket + * replacement (auth-token/user change): the tracked sessions belong to the old connection and + * will never emit another terminal event on this one. + */ + const reset = (): void => { + clearAll(); + }; + + return { onFinalImageLoaded, onItemStarted, onQueueCleared, onTerminal, recordProgress, reset }; +};