From 91ef24e15cb74ddef40cfc24281485a691b801ec Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Fri, 17 Nov 2023 07:32:04 +1100 Subject: [PATCH] fix(nodes,ui): fix missed/canvas temp images in gallery Resolves two bugs introduced in #5106: 1. Linear UI images sometimes didn't make it to the gallery. This was a race condition. The VAE decode nodes were handled by the socketInvocationComplete listener. At that moment, the image was marked as intermediate. Immediately after this node was handled, a LinearUIOutputInvocation, introduced in #5106, was handled by socketInvocationComplete. This node internally sets changed the image to not intermediate. During the handling of that socketInvocationComplete, RTK Query would sometimes use its cache instead of retrieving the image DTO again. The result is that the UI never got the message that the image was not intermediate, so it wasn't added to the gallery. This is resolved by refactoring the socketInvocationComplete listener. We now skip the gallery processing for linear UI events, except for the LinearUIOutputInvocation. Images now always make it to the gallery, and network requests to get image DTOs are substantially reduced. 2. Canvas temp images always went into the gallery The LinearUIOutputInvocation was always setting its image's is_intermediate to false. This included all canvas images and resulted in all canvas temp images going to gallery. This is resolved by making LinearUIOutputInvocation set is_intermediate based on `self.is_intermediate`. The behaviour now more or less mirroring the behaviour of is_intermediate on other image-outputting nodes, except it doesn't save the image again - only changes it. One extra minor change - LinearUIOutputInvocation only changes is_intermediate if it differs from the image's current setting. Very minor optimisation. --- invokeai/app/invocations/image.py | 7 +++-- .../socketio/socketInvocationComplete.ts | 28 +++++++++++++------ .../graphBuilders/buildAdHocUpscaleGraph.ts | 2 +- .../nodes/util/graphBuilders/constants.ts | 28 ++++++++++++++++++- 4 files changed, 53 insertions(+), 12 deletions(-) diff --git a/invokeai/app/invocations/image.py b/invokeai/app/invocations/image.py index fb59fdf7fd..87e8392402 100644 --- a/invokeai/app/invocations/image.py +++ b/invokeai/app/invocations/image.py @@ -1024,7 +1024,7 @@ class SaveImageInvocation(BaseInvocation, WithWorkflow, WithMetadata): title="Linear UI Image Output", tags=["primitives", "image"], category="primitives", - version="1.0.0", + version="1.0.1", use_cache=False, ) class LinearUIOutputInvocation(BaseInvocation, WithWorkflow, WithMetadata): @@ -1039,7 +1039,10 @@ class LinearUIOutputInvocation(BaseInvocation, WithWorkflow, WithMetadata): if self.board: context.services.board_images.add_image_to_board(self.board.board_id, self.image.image_name) - context.services.images.update(self.image.image_name, changes=ImageRecordChanges(is_intermediate=False)) + if image_dto.is_intermediate != self.is_intermediate: + context.services.images.update( + self.image.image_name, changes=ImageRecordChanges(is_intermediate=self.is_intermediate) + ) return ImageOutput( image=ImageField(image_name=self.image.image_name), diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketInvocationComplete.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketInvocationComplete.ts index 007a829bf9..cfd69ce9bc 100644 --- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketInvocationComplete.ts +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketInvocationComplete.ts @@ -7,7 +7,10 @@ import { imageSelected, } from 'features/gallery/store/gallerySlice'; import { IMAGE_CATEGORIES } from 'features/gallery/store/types'; -import { CANVAS_OUTPUT } from 'features/nodes/util/graphBuilders/constants'; +import { + LINEAR_UI_OUTPUT, + nodeIDDenyList, +} from 'features/nodes/util/graphBuilders/constants'; import { boardsApi } from 'services/api/endpoints/boards'; import { imagesApi } from 'services/api/endpoints/images'; import { isImageOutput } from 'services/api/guards'; @@ -19,7 +22,7 @@ import { import { startAppListening } from '../..'; // These nodes output an image, but do not actually *save* an image, so we don't want to handle the gallery logic on them -const nodeDenylist = ['load_image', 'image']; +const nodeTypeDenylist = ['load_image', 'image']; export const addInvocationCompleteEventListener = () => { startAppListening({ @@ -32,22 +35,31 @@ export const addInvocationCompleteEventListener = () => { `Invocation complete (${action.payload.data.node.type})` ); - const { result, node, queue_batch_id } = data; + const { result, node, queue_batch_id, source_node_id } = data; // This complete event has an associated image output - if (isImageOutput(result) && !nodeDenylist.includes(node.type)) { + if ( + isImageOutput(result) && + !nodeTypeDenylist.includes(node.type) && + !nodeIDDenyList.includes(source_node_id) + ) { const { image_name } = result.image; const { canvas, gallery } = getState(); // This populates the `getImageDTO` cache - const imageDTO = await dispatch( - imagesApi.endpoints.getImageDTO.initiate(image_name) - ).unwrap(); + const imageDTORequest = dispatch( + imagesApi.endpoints.getImageDTO.initiate(image_name, { + forceRefetch: true, + }) + ); + + const imageDTO = await imageDTORequest.unwrap(); + imageDTORequest.unsubscribe(); // Add canvas images to the staging area if ( canvas.batchIds.includes(queue_batch_id) && - [CANVAS_OUTPUT].includes(data.source_node_id) + [LINEAR_UI_OUTPUT].includes(data.source_node_id) ) { dispatch(addImageToStagingArea(imageDTO)); } diff --git a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildAdHocUpscaleGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildAdHocUpscaleGraph.ts index 73551202b8..8331c81eb3 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildAdHocUpscaleGraph.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildAdHocUpscaleGraph.ts @@ -6,7 +6,7 @@ import { Graph, LinearUIOutputInvocation, } from 'services/api/types'; -import { REALESRGAN as ESRGAN, LINEAR_UI_OUTPUT } from './constants'; +import { ESRGAN, LINEAR_UI_OUTPUT } from './constants'; import { addCoreMetadataNode, upsertMetadata } from './metadata'; type Arg = { diff --git a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/constants.ts b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/constants.ts index 3839822133..c2d6ac041c 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/constants.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/constants.ts @@ -67,7 +67,7 @@ export const BATCH_PROMPT = 'batch_prompt'; export const BATCH_STYLE_PROMPT = 'batch_style_prompt'; export const METADATA_COLLECT = 'metadata_collect'; export const MERGE_METADATA = 'merge_metadata'; -export const REALESRGAN = 'esrgan'; +export const ESRGAN = 'esrgan'; export const DIVIDE = 'divide'; export const SCALE = 'scale_image'; export const SDXL_MODEL_LOADER = 'sdxl_model_loader'; @@ -82,6 +82,32 @@ export const SDXL_REFINER_INPAINT_CREATE_MASK = 'refiner_inpaint_create_mask'; export const SEAMLESS = 'seamless'; export const SDXL_REFINER_SEAMLESS = 'refiner_seamless'; +// these image-outputting nodes are from the linear UI and we should not handle the gallery logic on them +// instead, we wait for LINEAR_UI_OUTPUT node, and handle it like any other image-outputting node +export const nodeIDDenyList = [ + CANVAS_OUTPUT, + LATENTS_TO_IMAGE, + LATENTS_TO_IMAGE_HRF_HR, + NSFW_CHECKER, + WATERMARKER, + ESRGAN, + ESRGAN_HRF, + RESIZE_HRF, + LATENTS_TO_IMAGE_HRF_LR, + IMG2IMG_RESIZE, + INPAINT_IMAGE, + SCALED_INPAINT_IMAGE, + INPAINT_IMAGE_RESIZE_UP, + INPAINT_IMAGE_RESIZE_DOWN, + INPAINT_INFILL, + INPAINT_INFILL_RESIZE_DOWN, + INPAINT_FINAL_IMAGE, + INPAINT_CREATE_MASK, + INPAINT_MASK, + PASTE_IMAGE, + SCALE, +]; + // friendly graph ids export const TEXT_TO_IMAGE_GRAPH = 'text_to_image_graph'; export const IMAGE_TO_IMAGE_GRAPH = 'image_to_image_graph';