From 2f26657c17dfc4241e4244a2d00dc4f86bdca1b2 Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Tue, 10 Jun 2025 12:16:33 +1000 Subject: [PATCH] feat(ui): move canvas-specific staging subscriptions to CanvasStagingAreaModule --- .../SimpleSession/StagingAreaItemsList.tsx | 27 +--------- .../components/SimpleSession/context.tsx | 54 +++++++++++++++---- .../konva/CanvasStagingAreaModule.ts | 27 +++++++++- 3 files changed, 71 insertions(+), 37 deletions(-) diff --git a/invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/StagingAreaItemsList.tsx b/invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/StagingAreaItemsList.tsx index c195fff0b1..36c6d21ec8 100644 --- a/invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/StagingAreaItemsList.tsx +++ b/invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/StagingAreaItemsList.tsx @@ -4,9 +4,7 @@ import { useStore } from '@nanostores/react'; import ScrollableContent from 'common/components/OverlayScrollbars/ScrollableContent'; import { useCanvasSessionContext } from 'features/controlLayers/components/SimpleSession/context'; import { QueueItemPreviewMini } from 'features/controlLayers/components/SimpleSession/QueueItemPreviewMini'; -import { getOutputImageName } from 'features/controlLayers/components/SimpleSession/shared'; import { useCanvasManagerSafe } from 'features/controlLayers/contexts/CanvasManagerProviderGate'; -import { effect } from 'nanostores'; import { memo, useEffect } from 'react'; export const StagingAreaItemsList = memo(() => { @@ -20,29 +18,8 @@ export const StagingAreaItemsList = memo(() => { return; } - return effect([ctx.$selectedItem, ctx.$progressData], (selectedItem, progressData) => { - if (!selectedItem) { - canvasManager.stagingArea.$imageSrc.set(null); - return; - } - - const outputImageName = getOutputImageName(selectedItem); - - if (outputImageName) { - canvasManager.stagingArea.$imageSrc.set({ type: 'imageName', data: outputImageName }); - return; - } - - const data = progressData[selectedItem.item_id]; - - if (data?.progressImage) { - canvasManager.stagingArea.$imageSrc.set({ type: 'dataURL', data: data.progressImage.dataURL }); - return; - } - - canvasManager.stagingArea.$imageSrc.set(null); - }); - }, [canvasManager, ctx.$progressData, ctx.$selectedItem]); + return canvasManager.stagingArea.connectToSession(ctx.$selectedItemId, ctx.$progressData); + }, [canvasManager, ctx.$progressData, ctx.$selectedItemId]); return ( diff --git a/invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/context.tsx b/invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/context.tsx index daf843f947..172a8188ce 100644 --- a/invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/context.tsx +++ b/invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/context.tsx @@ -13,18 +13,26 @@ import type { S } from 'services/api/types'; import { $socket } from 'services/events/stores'; import { assert } from 'tsafe'; -type ProgressData = { +export type ProgressData = { itemId: number; progressEvent: S['InvocationProgressEvent'] | null; progressImage: ProgressImage | null; + outputImageName: string | null; }; +const getInitialProgressData = (itemId: number): ProgressData => ({ + itemId, + progressEvent: null, + progressImage: null, + outputImageName: null, +}); + export const useProgressData = ( $progressData: WritableAtom>, itemId: number ): ProgressData => { const [value, setValue] = useState(() => { - return $progressData.get()[itemId] ?? { itemId, progressEvent: null, progressImage: null }; + return $progressData.get()[itemId] ?? getInitialProgressData(itemId); }); useEffect(() => { const unsub = $progressData.subscribe((data) => { @@ -62,6 +70,7 @@ const setProgress = ($progressData: WritableAtom>, itemId: data.item_id, progressEvent: data, progressImage: data.image ?? null, + outputImageName: null, }, }); } @@ -315,18 +324,45 @@ export const CanvasSessionContextProvider = memo( const progressData = $progressData.get(); const toDelete: number[] = []; - const toClear: number[] = []; + const toUpdate: ProgressData[] = []; for (const datum of Object.values(progressData)) { const item = items.find(({ item_id }) => item_id === datum.itemId); if (!item) { toDelete.push(datum.itemId); } else if (item.status === 'canceled' || item.status === 'failed') { - toClear.push(datum.itemId); + toUpdate[datum.itemId] = { + ...datum, + progressEvent: null, + progressImage: null, + outputImageName: null, + }; } } - if (toDelete.length === 0) { + for (const item of items) { + const datum = progressData[item.item_id]; + + if (datum) { + if (datum.outputImageName) { + continue; + } + const outputImageName = getOutputImageName(item); + if (!outputImageName) { + continue; + } + toUpdate.push({ + ...datum, + outputImageName, + }); + } else { + const _datum = getInitialProgressData(item.item_id); + _datum.outputImageName = getOutputImageName(item); + toUpdate.push(_datum); + } + } + + if (toDelete.length === 0 && toUpdate.length === 0) { return; } @@ -336,12 +372,8 @@ export const CanvasSessionContextProvider = memo( delete newProgressData[itemId]; } - for (const itemId of toClear) { - const current = newProgressData[itemId]; - if (current) { - current.progressEvent = null; - current.progressImage = null; - } + for (const datum of toUpdate) { + newProgressData[datum.itemId] = datum; } $progressData.set(newProgressData); diff --git a/invokeai/frontend/web/src/features/controlLayers/konva/CanvasStagingAreaModule.ts b/invokeai/frontend/web/src/features/controlLayers/konva/CanvasStagingAreaModule.ts index 52cbc54611..fa3ca84ec6 100644 --- a/invokeai/frontend/web/src/features/controlLayers/konva/CanvasStagingAreaModule.ts +++ b/invokeai/frontend/web/src/features/controlLayers/konva/CanvasStagingAreaModule.ts @@ -1,4 +1,5 @@ import { Mutex } from 'async-mutex'; +import type { ProgressData } from 'features/controlLayers/components/SimpleSession/context'; import type { CanvasManager } from 'features/controlLayers/konva/CanvasManager'; import { CanvasModuleBase } from 'features/controlLayers/konva/CanvasModuleBase'; import { CanvasObjectImage } from 'features/controlLayers/konva/CanvasObject/CanvasObjectImage'; @@ -6,7 +7,8 @@ import { getPrefixedId } from 'features/controlLayers/konva/util'; import { selectIsStaging } from 'features/controlLayers/store/canvasStagingAreaSlice'; import type { CanvasImageState } from 'features/controlLayers/store/types'; import Konva from 'konva'; -import { atom } from 'nanostores'; +import type { Atom, WritableAtom } from 'nanostores'; +import { atom, effect } from 'nanostores'; import type { Logger } from 'roarr'; type ImageNameSrc = { type: 'imageName'; data: string }; @@ -133,6 +135,29 @@ export class CanvasStagingAreaModule extends CanvasModuleBase { this.$isStaging.set(this.manager.stateApi.runSelector(selectIsStaging)); }; + connectToSession = ( + $selectedItemId: Atom, + $progressData: WritableAtom> + ) => + effect([$selectedItemId, $progressData], (selectedItemId, progressData) => { + if (!selectedItemId) { + this.$imageSrc.set(null); + return; + } + + const datum = progressData[selectedItemId]; + + if (datum?.outputImageName) { + this.$imageSrc.set({ type: 'imageName', data: datum.outputImageName }); + return; + } else if (datum?.progressImage) { + this.$imageSrc.set({ type: 'dataURL', data: datum.progressImage.dataURL }); + return; + } else { + this.$imageSrc.set(null); + } + }); + private _getImageFromSrc = ( { type, data }: ImageNameSrc | DataURLSrc, width: number,