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 bf219cd160..fe427a13c7 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 @@ -49,7 +49,6 @@ export const buildSD1Graph = async (state: RootState, manager: KonvaNodeManager) vae, positivePrompt, negativePrompt, - img2imgStrength, } = params; assert(model, 'No model found in state'); @@ -108,9 +107,6 @@ export const buildSD1Graph = async (state: RootState, manager: KonvaNodeManager) id: LATENTS_TO_IMAGE, fp32: vaePrecision === 'fp32', board: getBoardField(state), - // This is the terminal node and must always save to gallery. - is_intermediate: false, - use_cache: false, }); const vaeLoader = vae?.base === model.base @@ -121,8 +117,7 @@ export const buildSD1Graph = async (state: RootState, manager: KonvaNodeManager) }) : null; - let imageOutput: Invocation<'l2i'> | Invocation<'img_nsfw'> | Invocation<'img_watermark'> | Invocation<'img_resize'> = - l2i; + let imageOutput: Invocation<'l2i' | 'img_nsfw' | 'img_watermark' | 'img_resize' | 'canvas_paste_back'> = l2i; g.addEdge(modelLoader, 'unet', denoise, 'unet'); g.addEdge(modelLoader, 'clip', clipSkip, 'clip'); @@ -165,51 +160,173 @@ export const buildSD1Graph = async (state: RootState, manager: KonvaNodeManager) if (generationMode === 'txt2img') { if (!isEqual(scaledSize, originalSize)) { - // We are using scaled bbox and need to resize the output image back to the original size. - imageOutput = g.addNode({ - id: 'img_resize', + // We need to resize the output image back to the original size + const resizeImageToOriginalSize = g.addNode({ + id: 'resize_image_to_original_size', type: 'img_resize', ...originalSize, - is_intermediate: false, - use_cache: false, }); - g.addEdge(l2i, 'image', imageOutput, 'image'); + g.addEdge(l2i, 'image', resizeImageToOriginalSize, 'image'); + + // This is the new output node + imageOutput = resizeImageToOriginalSize; } } else if (generationMode === 'img2img') { - const { image_name } = await manager.util.getImageSourceImage({ - bbox: pick(bbox, ['x', 'y', 'width', 'height']), + denoise.denoising_start = 1 - params.img2imgStrength; + + const cropBbox = pick(bbox, ['x', 'y', 'width', 'height']); + const initialImage = await manager.util.getImageSourceImage({ + bbox: cropBbox, preview: true, }); - denoise.denoising_start = 1 - img2imgStrength; - if (!isEqual(scaledSize, originalSize)) { - // We are using scaled bbox and need to resize the output image back to the original size. - const initialImageResize = g.addNode({ - id: 'initial_image_resize', + // Resize the initial image to the scaled size, denoise, then resize back to the original size + const resizeImageToScaledSize = g.addNode({ + id: 'initial_image_resize_in', type: 'img_resize', + image: { image_name: initialImage.image_name }, ...scaledSize, - image: { image_name }, }); const i2l = g.addNode({ id: 'i2l', type: 'i2l' }); - - g.addEdge(vaeSource, 'vae', i2l, 'vae'); - g.addEdge(initialImageResize, 'image', i2l, 'image'); - g.addEdge(i2l, 'latents', denoise, 'latents'); - - imageOutput = g.addNode({ - id: 'img_resize', + const resizeImageToOriginalSize = g.addNode({ + id: 'initial_image_resize_out', type: 'img_resize', ...originalSize, - is_intermediate: false, - use_cache: false, }); - g.addEdge(l2i, 'image', imageOutput, 'image'); + + g.addEdge(vaeSource, 'vae', i2l, 'vae'); + g.addEdge(resizeImageToScaledSize, 'image', i2l, 'image'); + g.addEdge(i2l, 'latents', denoise, 'latents'); + g.addEdge(l2i, 'image', resizeImageToOriginalSize, 'image'); + + // This is the new output node + imageOutput = resizeImageToOriginalSize; } else { - const i2l = g.addNode({ id: 'i2l', type: 'i2l', image: { image_name } }); + // No need to resize, just denoise + const i2l = g.addNode({ id: 'i2l', type: 'i2l', image: { image_name: initialImage.image_name } }); g.addEdge(vaeSource, 'vae', i2l, 'vae'); g.addEdge(i2l, 'latents', denoise, 'latents'); } + } else if (generationMode === 'inpaint') { + denoise.denoising_start = 1 - params.img2imgStrength; + + const { compositing } = state.canvasV2; + + const cropBbox = pick(bbox, ['x', 'y', 'width', 'height']); + const initialImage = await manager.util.getImageSourceImage({ + bbox: cropBbox, + preview: true, + }); + const maskImage = await manager.util.getInpaintMaskImage({ + bbox: cropBbox, + preview: true, + }); + + if (!isEqual(scaledSize, originalSize)) { + // Scale before processing requires some resizing + const i2l = g.addNode({ id: 'i2l', type: 'i2l' }); + const resizeImageToScaledSize = g.addNode({ + id: 'resize_image_to_scaled_size', + type: 'img_resize', + image: { image_name: initialImage.image_name }, + ...scaledSize, + }); + const alphaToMask = g.addNode({ + id: 'alpha_to_mask', + type: 'tomask', + image: { image_name: maskImage.image_name }, + invert: true, + }); + const resizeMaskToScaledSize = g.addNode({ + id: 'resize_mask_to_scaled_size', + type: 'img_resize', + ...scaledSize, + }); + const resizeImageToOriginalSize = g.addNode({ + id: 'resize_image_to_original_size', + type: 'img_resize', + ...originalSize, + }); + const resizeMaskToOriginalSize = g.addNode({ + id: 'resize_mask_to_original_size', + type: 'img_resize', + ...originalSize, + }); + const createGradientMask = g.addNode({ + id: 'create_gradient_mask', + type: 'create_gradient_mask', + coherence_mode: compositing.canvasCoherenceMode, + minimum_denoise: compositing.canvasCoherenceMinDenoise, + edge_radius: compositing.canvasCoherenceEdgeSize, + fp32: vaePrecision === 'fp32', + }); + const canvasPasteBack = g.addNode({ + id: 'canvas_paste_back', + type: 'canvas_paste_back', + board: getBoardField(state), + mask_blur: compositing.maskBlur, + source_image: { image_name: initialImage.image_name }, + }); + + // Resize initial image and mask to scaled size, feed into to gradient mask + g.addEdge(alphaToMask, 'image', resizeMaskToScaledSize, 'image'); + g.addEdge(resizeImageToScaledSize, 'image', i2l, 'image'); + g.addEdge(i2l, 'latents', denoise, 'latents'); + g.addEdge(vaeSource, 'vae', i2l, 'vae'); + + g.addEdge(vaeSource, 'vae', createGradientMask, 'vae'); + g.addEdge(modelLoader, 'unet', createGradientMask, 'unet'); + g.addEdge(resizeImageToScaledSize, 'image', createGradientMask, 'image'); + g.addEdge(resizeMaskToScaledSize, 'image', createGradientMask, 'mask'); + + g.addEdge(createGradientMask, 'denoise_mask', denoise, 'denoise_mask'); + + // After denoising, resize the image and mask back to original size + g.addEdge(l2i, 'image', resizeImageToOriginalSize, 'image'); + g.addEdge(createGradientMask, 'expanded_mask_area', resizeMaskToOriginalSize, 'image'); + + // Finally, paste the generated masked image back onto the original image + g.addEdge(resizeImageToOriginalSize, 'image', canvasPasteBack, 'target_image'); + g.addEdge(resizeMaskToOriginalSize, 'image', canvasPasteBack, 'mask'); + + imageOutput = canvasPasteBack; + } else { + // No scale before processing, much simpler + const i2l = g.addNode({ id: 'i2l', type: 'i2l', image: { image_name: initialImage.image_name } }); + const alphaToMask = g.addNode({ + id: 'alpha_to_mask', + type: 'tomask', + image: { image_name: maskImage.image_name }, + invert: true, + }); + const createGradientMask = g.addNode({ + id: 'create_gradient_mask', + type: 'create_gradient_mask', + coherence_mode: compositing.canvasCoherenceMode, + minimum_denoise: compositing.canvasCoherenceMinDenoise, + edge_radius: compositing.canvasCoherenceEdgeSize, + fp32: vaePrecision === 'fp32', + image: { image_name: initialImage.image_name }, + }); + const canvasPasteBack = g.addNode({ + id: 'canvas_paste_back', + type: 'canvas_paste_back', + board: getBoardField(state), + mask_blur: compositing.maskBlur, + source_image: { image_name: initialImage.image_name }, + mask: { image_name: maskImage.image_name }, + }); + g.addEdge(alphaToMask, 'image', createGradientMask, 'mask'); + g.addEdge(i2l, 'latents', denoise, 'latents'); + g.addEdge(vaeSource, 'vae', i2l, 'vae'); + g.addEdge(vaeSource, 'vae', createGradientMask, 'vae'); + g.addEdge(modelLoader, 'unet', createGradientMask, 'unet'); + g.addEdge(createGradientMask, 'denoise_mask', denoise, 'denoise_mask'); + g.addEdge(l2i, 'image', canvasPasteBack, 'target_image'); + + imageOutput = canvasPasteBack; + } } const _addedCAs = addControlAdapters(state.canvasV2.controlAdapters.entities, g, denoise, modelConfig.base); @@ -241,6 +358,10 @@ export const buildSD1Graph = async (state: RootState, manager: KonvaNodeManager) imageOutput = addWatermarker(g, imageOutput); } + // This is the terminal node and must always save to gallery. + imageOutput.is_intermediate = false; + imageOutput.use_cache = false; + g.setMetadataReceivingNode(imageOutput); return g.getGraph(); };