diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/addInpaint.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/addInpaint.ts index ef0aed835b..35f512a1c9 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/addInpaint.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/addInpaint.ts @@ -1,3 +1,4 @@ +import type { RootState } from 'app/store/store'; import type { CanvasManager } from 'features/controlLayers/konva/CanvasManager'; import { getPrefixedId } from 'features/controlLayers/konva/util'; import type { CanvasV2State, Dimensions } from 'features/controlLayers/store/types'; @@ -7,6 +8,7 @@ import { isEqual } from 'lodash-es'; import type { Invocation } from 'services/api/types'; export const addInpaint = async ( + state: RootState, g: Graph, manager: CanvasManager, l2i: Invocation<'l2i'>, @@ -22,6 +24,7 @@ export const addInpaint = async ( ): Promise> => { denoise.denoising_start = denoising_start; + const mode = state.canvasV2.session.mode; const initialImage = await manager.compositor.getCompositeRasterLayerImageDTO(bbox.rect); const maskImage = await manager.compositor.getCompositeInpaintMaskImageDTO(bbox.rect); @@ -87,9 +90,13 @@ export const addInpaint = async ( g.addEdge(createGradientMask, 'expanded_mask_area', resizeMaskToOriginalSize, 'image'); // Finally, paste the generated masked image back onto the original image - g.addEdge(resizeImageToOriginalSize, 'image', canvasPasteBack, 'image'); + g.addEdge(resizeImageToOriginalSize, 'image', canvasPasteBack, 'generated_image'); g.addEdge(resizeMaskToOriginalSize, 'image', canvasPasteBack, 'mask'); + if (mode === 'generate') { + canvasPasteBack.source_image = { image_name: initialImage.image_name }; + } + return canvasPasteBack; } else { // No scale before processing, much simpler @@ -114,6 +121,7 @@ export const addInpaint = async ( type: 'canvas_v2_mask_and_crop', mask_blur: compositing.maskBlur, }); + g.addEdge(alphaToMask, 'image', createGradientMask, 'mask'); g.addEdge(i2l, 'latents', denoise, 'latents'); g.addEdge(vaeSource, 'vae', i2l, 'vae'); @@ -122,7 +130,11 @@ export const addInpaint = async ( g.addEdge(createGradientMask, 'denoise_mask', denoise, 'denoise_mask'); g.addEdge(createGradientMask, 'expanded_mask_area', canvasPasteBack, 'mask'); - g.addEdge(l2i, 'image', canvasPasteBack, 'image'); + g.addEdge(l2i, 'image', canvasPasteBack, 'generated_image'); + + if (mode === 'generate') { + canvasPasteBack.source_image = { image_name: initialImage.image_name }; + } return canvasPasteBack; } diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/addOutpaint.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/addOutpaint.ts index b2d5b4a18e..55fd2eac4e 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/addOutpaint.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/addOutpaint.ts @@ -1,3 +1,4 @@ +import type { RootState } from 'app/store/store'; import type { CanvasManager } from 'features/controlLayers/konva/CanvasManager'; import { getPrefixedId } from 'features/controlLayers/konva/util'; import type { CanvasV2State, Dimensions } from 'features/controlLayers/store/types'; @@ -8,6 +9,7 @@ import { isEqual } from 'lodash-es'; import type { Invocation } from 'services/api/types'; export const addOutpaint = async ( + state: RootState, g: Graph, manager: CanvasManager, l2i: Invocation<'l2i'>, @@ -23,6 +25,7 @@ export const addOutpaint = async ( ): Promise> => { denoise.denoising_start = denoising_start; + const mode = state.canvasV2.session.mode; const initialImage = await manager.compositor.getCompositeRasterLayerImageDTO(bbox.rect); const maskImage = await manager.compositor.getCompositeInpaintMaskImageDTO(bbox.rect); const infill = getInfill(g, compositing); @@ -111,9 +114,13 @@ export const addOutpaint = async ( g.addEdge(createGradientMask, 'expanded_mask_area', resizeOutputMaskToOriginalSize, 'image'); // Finally, paste the generated masked image back onto the original image - g.addEdge(resizeOutputImageToOriginalSize, 'image', canvasPasteBack, 'image'); + g.addEdge(resizeOutputImageToOriginalSize, 'image', canvasPasteBack, 'generated_image'); g.addEdge(resizeOutputMaskToOriginalSize, 'image', canvasPasteBack, 'mask'); + if (mode === 'generate') { + canvasPasteBack.source_image = { image_name: initialImage.image_name }; + } + return canvasPasteBack; } else { infill.image = { image_name: initialImage.image_name }; @@ -158,7 +165,11 @@ export const addOutpaint = async ( g.addEdge(modelLoader, 'unet', createGradientMask, 'unet'); g.addEdge(createGradientMask, 'denoise_mask', denoise, 'denoise_mask'); g.addEdge(createGradientMask, 'expanded_mask_area', canvasPasteBack, 'mask'); - g.addEdge(l2i, 'image', canvasPasteBack, 'image'); + g.addEdge(l2i, 'image', canvasPasteBack, 'generated_image'); + + if (mode === 'generate') { + canvasPasteBack.source_image = { image_name: initialImage.image_name }; + } return canvasPasteBack; } diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSD1Graph.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSD1Graph.ts index b0aad315ca..237acfca13 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSD1Graph.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSD1Graph.ts @@ -173,6 +173,7 @@ export const buildSD1Graph = async ( } else if (generationMode === 'inpaint') { const { compositing } = state.canvasV2; canvasOutput = await addInpaint( + state, g, manager, l2i, @@ -189,6 +190,7 @@ export const buildSD1Graph = async ( } else if (generationMode === 'outpaint') { const { compositing } = state.canvasV2; canvasOutput = await addOutpaint( + state, g, manager, l2i, diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSDXLGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSDXLGraph.ts index 601e9671a7..6e8c8a7248 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSDXLGraph.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSDXLGraph.ts @@ -176,6 +176,7 @@ export const buildSDXLGraph = async ( } else if (generationMode === 'inpaint') { const { compositing } = state.canvasV2; canvasOutput = await addInpaint( + state, g, manager, l2i, @@ -192,6 +193,7 @@ export const buildSDXLGraph = async ( } else if (generationMode === 'outpaint') { const { compositing } = state.canvasV2; canvasOutput = await addOutpaint( + state, g, manager, l2i,