diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketGeneratorProgress.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketGeneratorProgress.ts index 08ad830ba4..14af8bdbf4 100644 --- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketGeneratorProgress.ts +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketGeneratorProgress.ts @@ -2,6 +2,7 @@ import { logger } from 'app/logging/logger'; import type { AppStartListening } from 'app/store/middleware/listenerMiddleware'; import { deepClone } from 'common/util/deepClone'; import { parseify } from 'common/util/serialize'; +import { $lastProgressEvent } from 'features/controlLayers/store/canvasV2Slice'; import { $nodeExecutionStates, upsertExecutionState } from 'features/nodes/hooks/useExecutionState'; import { zNodeStatus } from 'features/nodes/types/invocation'; import { socketGeneratorProgress } from 'services/events/actions'; @@ -11,9 +12,9 @@ const log = logger('socketio'); export const addGeneratorProgressEventListener = (startAppListening: AppStartListening) => { startAppListening({ actionCreator: socketGeneratorProgress, - effect: (action) => { + effect: (action, { getState }) => { log.trace(parseify(action.payload), `Generator progress`); - const { invocation_source_id, step, total_steps, progress_image } = action.payload.data; + const { invocation_source_id, step, total_steps, progress_image, batch_id } = action.payload.data; const nes = deepClone($nodeExecutionStates.get()[invocation_source_id]); if (nes) { nes.status = zNodeStatus.enum.IN_PROGRESS; @@ -21,6 +22,11 @@ export const addGeneratorProgressEventListener = (startAppListening: AppStartLis nes.progressImage = progress_image ?? null; upsertExecutionState(nes.nodeId, nes); } + + const isCanvasQueueItem = getState().canvasV2.stagingArea?.batchIds.includes(batch_id); + if (isCanvasQueueItem) { + $lastProgressEvent.set(action.payload.data); + } }, }); }; diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketQueueItemStatusChanged.tsx b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketQueueItemStatusChanged.tsx index 8a83609b3c..2d117ef140 100644 --- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketQueueItemStatusChanged.tsx +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketQueueItemStatusChanged.tsx @@ -1,6 +1,7 @@ import { logger } from 'app/logging/logger'; import type { AppStartListening } from 'app/store/middleware/listenerMiddleware'; import { deepClone } from 'common/util/deepClone'; +import { $lastProgressEvent } from 'features/controlLayers/store/canvasV2Slice'; import { $nodeExecutionStates } from 'features/nodes/hooks/useExecutionState'; import { zNodeStatus } from 'features/nodes/types/invocation'; import ErrorToastDescription, { getTitleFromErrorType } from 'features/toast/ErrorToastDescription'; @@ -28,10 +29,13 @@ export const addSocketQueueItemStatusChangedEventListener = (startAppListening: error_type, error_message, error_traceback, + batch_id, } = action.payload.data; log.debug(action.payload, `Queue item ${item_id} status updated: ${status}`); + const isCanvasQueueItem = getState().canvasV2.stagingArea?.batchIds.includes(batch_id); + // Update this specific queue item in the list of queue items (this is the queue item DTO, without the session) dispatch( queueApi.util.updateQueryData('listQueueItems', undefined, (draft) => { @@ -92,6 +96,9 @@ export const addSocketQueueItemStatusChangedEventListener = (startAppListening: } else if (status === 'failed' && error_type) { const isLocal = getState().config.isLocal ?? true; const sessionId = session_id; + if (isCanvasQueueItem) { + $lastProgressEvent.set(null); + } toast({ id: `INVOCATION_ERROR_${error_type}`, @@ -108,6 +115,10 @@ export const addSocketQueueItemStatusChangedEventListener = (startAppListening: /> ), }); + } else if (status === 'completed' && isCanvasQueueItem) { + $lastProgressEvent.set(null); + } else if (status === 'canceled' && isCanvasQueueItem) { + $lastProgressEvent.set(null); } }, }); diff --git a/invokeai/frontend/web/src/features/controlLayers/konva/nodeManager.ts b/invokeai/frontend/web/src/features/controlLayers/konva/nodeManager.ts index 6dd0a591bd..941d34a106 100644 --- a/invokeai/frontend/web/src/features/controlLayers/konva/nodeManager.ts +++ b/invokeai/frontend/web/src/features/controlLayers/konva/nodeManager.ts @@ -22,6 +22,7 @@ import type { Vector2d } from 'konva/lib/types'; import { atom } from 'nanostores'; import { getImageDTO as defaultGetImageDTO, uploadImage as defaultUploadImage } from 'services/api/endpoints/images'; import type { ImageCategory, ImageDTO } from 'services/api/types'; +import type { InvocationDenoiseProgressEvent } from 'services/events/types'; import { assert } from 'tsafe'; import { CanvasBbox } from './renderers/bbox'; @@ -74,6 +75,7 @@ export type StateApi = { getRegionsState: () => CanvasV2State['regions']; getInpaintMaskState: () => CanvasV2State['inpaintMask']; getStagingAreaState: () => CanvasV2State['stagingArea']; + getLastProgressEvent: () => InvocationDenoiseProgressEvent | null; onInpaintMaskImageCached: (imageDTO: ImageDTO) => void; onRegionMaskImageCached: (id: string, imageDTO: ImageDTO) => void; onLayerImageCached: (imageDTO: ImageDTO) => void; @@ -276,7 +278,11 @@ export class KonvaNodeManager { } renderStagingArea() { - this.preview.stagingArea.render(this.stateApi.getStagingAreaState(), this.stateApi.getShouldShowStagedImage()); + this.preview.stagingArea.render( + this.stateApi.getStagingAreaState(), + this.stateApi.getShouldShowStagedImage(), + this.stateApi.getLastProgressEvent() + ); } fitDocument() { diff --git a/invokeai/frontend/web/src/features/controlLayers/konva/renderers/controlAdapters.ts b/invokeai/frontend/web/src/features/controlLayers/konva/renderers/controlAdapters.ts index 13f2b3a899..ea696cf2dc 100644 --- a/invokeai/frontend/web/src/features/controlLayers/konva/renderers/controlAdapters.ts +++ b/invokeai/frontend/web/src/features/controlLayers/konva/renderers/controlAdapters.ts @@ -33,7 +33,7 @@ export class CanvasControlAdapter { const imageObject = entity.processedImageObject ?? entity.imageObject; if (!imageObject) { if (this.image) { - this.image.destroy(); + this.image.konvaImageGroup.visible(false); } return; } diff --git a/invokeai/frontend/web/src/features/controlLayers/konva/renderers/objects.ts b/invokeai/frontend/web/src/features/controlLayers/konva/renderers/objects.ts index 5a0c90b805..1c1230e0c5 100644 --- a/invokeai/frontend/web/src/features/controlLayers/konva/renderers/objects.ts +++ b/invokeai/frontend/web/src/features/controlLayers/konva/renderers/objects.ts @@ -173,6 +173,7 @@ export class KonvaImage { this.konvaPlaceholderGroup.visible(true); this.konvaPlaceholderText.text(t('common.loadingImage', 'Loading Image')); } + this.konvaImageGroup.visible(true); if (onLoading) { onLoading(); } @@ -194,6 +195,8 @@ export class KonvaImage { this.isLoading = false; this.isError = false; this.konvaPlaceholderGroup.visible(false); + this.konvaImageGroup.visible(true); + if (onLoad) { onLoad(this.konvaImage); } @@ -204,6 +207,8 @@ export class KonvaImage { this.isError = true; this.konvaPlaceholderGroup.visible(true); this.konvaPlaceholderText.text(t('common.imageFailedToLoad', 'Image Failed to Load')); + this.konvaImageGroup.visible(true); + if (onError) { onError(); } @@ -237,3 +242,62 @@ export class KonvaImage { this.konvaImageGroup.destroy(); } } + +export class KonvaProgressImage { + id: string; + progressImageId: string | null; + konvaImageGroup: Konva.Group; + konvaImage: Konva.Image | null; // The image is loaded asynchronously, so it may not be available immediately + isLoading: boolean; + isError: boolean; + + constructor(arg: { id: string }) { + const { id } = arg; + this.konvaImageGroup = new Konva.Group({ id, listening: false }); + + this.id = id; + this.progressImageId = null; + this.konvaImage = null; + this.isLoading = false; + this.isError = false; + } + + async updateImageSource( + progressImageId: string, + dataURL: string, + x: number, + y: number, + width: number, + height: number + ) { + const imageEl = new Image(); + imageEl.onload = () => { + if (this.konvaImage) { + this.konvaImage.setAttrs({ + image: imageEl, + x, + y, + width, + height, + }); + } else { + this.konvaImage = new Konva.Image({ + id: this.id, + listening: false, + image: imageEl, + x, + y, + width, + height, + }); + this.konvaImageGroup.add(this.konvaImage); + } + }; + imageEl.id = progressImageId; + imageEl.src = dataURL; + } + + destroy() { + this.konvaImageGroup.destroy(); + } +} diff --git a/invokeai/frontend/web/src/features/controlLayers/konva/renderers/renderer.ts b/invokeai/frontend/web/src/features/controlLayers/konva/renderers/renderer.ts index 5e5ee6b4d7..fa692558d8 100644 --- a/invokeai/frontend/web/src/features/controlLayers/konva/renderers/renderer.ts +++ b/invokeai/frontend/web/src/features/controlLayers/konva/renderers/renderer.ts @@ -7,6 +7,7 @@ import { setStageEventHandlers } from 'features/controlLayers/konva/events'; import { KonvaNodeManager, setNodeManager } from 'features/controlLayers/konva/nodeManager'; import { updateBboxes } from 'features/controlLayers/konva/renderers/entityBbox'; import { + $lastProgressEvent, $shouldShowStagedImage, $stageAttrs, bboxChanged, @@ -305,6 +306,7 @@ export const initializeRenderer = ( getInpaintMaskState, getStagingAreaState, getShouldShowStagedImage: $shouldShowStagedImage.get, + getLastProgressEvent: $lastProgressEvent.get, // Read-write state setTool, @@ -453,6 +455,11 @@ export const initializeRenderer = ( } }); + $lastProgressEvent.subscribe(() => { + logIfDebugging('Rendering staging area'); + manager.renderStagingArea(); + }); + logIfDebugging('First render of konva stage'); // On first render, the document should be fit to the stage. manager.renderDocumentSizeOverlay(); diff --git a/invokeai/frontend/web/src/features/controlLayers/konva/renderers/stagingArea.ts b/invokeai/frontend/web/src/features/controlLayers/konva/renderers/stagingArea.ts index 93c8cdb7d7..28b64f898b 100644 --- a/invokeai/frontend/web/src/features/controlLayers/konva/renderers/stagingArea.ts +++ b/invokeai/frontend/web/src/features/controlLayers/konva/renderers/stagingArea.ts @@ -1,34 +1,58 @@ -import { KonvaImage } from 'features/controlLayers/konva/renderers/objects'; +import { KonvaImage, KonvaProgressImage } from 'features/controlLayers/konva/renderers/objects'; import type { CanvasV2State } from 'features/controlLayers/store/types'; import Konva from 'konva'; +import type { InvocationDenoiseProgressEvent } from 'services/events/types'; import { assert } from 'tsafe'; export class CanvasStagingArea { group: Konva.Group; image: KonvaImage | null; + progressImage: KonvaProgressImage | null; constructor() { this.group = new Konva.Group({ listening: false }); this.image = null; + this.progressImage = null; } - async render(stagingArea: CanvasV2State['stagingArea'], shouldShowStagedImage: boolean) { - if (!stagingArea || stagingArea.selectedImageIndex === null) { - if (this.image) { - this.image.destroy(); - this.image = null; + async render( + stagingArea: CanvasV2State['stagingArea'], + shouldShowStagedImage: boolean, + lastProgressEvent: InvocationDenoiseProgressEvent | null + ) { + if (stagingArea && lastProgressEvent) { + const { invocation, step, progress_image } = lastProgressEvent; + const { dataURL } = progress_image; + const { x, y, width, height } = stagingArea.bbox; + const progressImageId = `${invocation.id}_${step}`; + if (this.progressImage) { + if ( + !this.progressImage.isLoading && + !this.progressImage.isError && + this.progressImage.progressImageId !== progressImageId + ) { + await this.progressImage.updateImageSource(progressImageId, dataURL, x, y, width, height); + this.image?.konvaImageGroup.visible(false); + this.progressImage.konvaImageGroup.visible(true); + } + } else { + this.progressImage = new KonvaProgressImage({ id: 'progress-image' }); + this.group.add(this.progressImage.konvaImageGroup); + await this.progressImage.updateImageSource(progressImageId, dataURL, x, y, width, height); + this.image?.konvaImageGroup.visible(false); + this.progressImage.konvaImageGroup.visible(true); } - return; - } - - if (stagingArea.selectedImageIndex !== null) { + } else if (stagingArea && stagingArea.selectedImageIndex !== null) { const imageDTO = stagingArea.images[stagingArea.selectedImageIndex]; assert(imageDTO, 'Image must exist'); if (this.image) { if (!this.image.isLoading && !this.image.isError && this.image.imageName !== imageDTO.image_name) { await this.image.updateImageSource(imageDTO.image_name); } + this.image.konvaImageGroup.x(stagingArea.bbox.x); + this.image.konvaImageGroup.y(stagingArea.bbox.y); this.image.konvaImageGroup.visible(shouldShowStagedImage); + this.progressImage?.konvaImageGroup.visible(false); } else { const { image_name, width, height } = imageDTO; this.image = new KonvaImage({ @@ -50,6 +74,14 @@ export class CanvasStagingArea { this.group.add(this.image.konvaImageGroup); await this.image.updateImageSource(imageDTO.image_name); this.image.konvaImageGroup.visible(shouldShowStagedImage); + this.progressImage?.konvaImageGroup.visible(false); + } + } else { + if (this.image) { + this.image.konvaImageGroup.visible(false); + } + if (this.progressImage) { + this.progressImage.konvaImageGroup.visible(false); } } } diff --git a/invokeai/frontend/web/src/features/controlLayers/store/canvasV2Slice.ts b/invokeai/frontend/web/src/features/controlLayers/store/canvasV2Slice.ts index ac4b385a25..02804d7bb4 100644 --- a/invokeai/frontend/web/src/features/controlLayers/store/canvasV2Slice.ts +++ b/invokeai/frontend/web/src/features/controlLayers/store/canvasV2Slice.ts @@ -18,6 +18,7 @@ import { toolReducers } from 'features/controlLayers/store/toolReducers'; import { initialAspectRatioState } from 'features/parameters/components/ImageSize/constants'; import type { AspectRatioState } from 'features/parameters/components/ImageSize/types'; import { atom } from 'nanostores'; +import type { InvocationDenoiseProgressEvent } from 'services/events/types'; import type { CanvasEntityIdentifier, CanvasV2State, StageAttrs } from './types'; import { RGBA_RED } from './types'; @@ -358,6 +359,7 @@ export const $stageAttrs = atom({ scale: 0, }); export const $shouldShowStagedImage = atom(true); +export const $lastProgressEvent = atom(null); export const canvasV2PersistConfig: PersistConfig = { name: canvasV2Slice.name,