From 9846229e52f11dd7e341885bc7201b14289eb4e9 Mon Sep 17 00:00:00 2001 From: Mary Hipp Date: Mon, 17 Mar 2025 14:57:50 -0400 Subject: [PATCH] build graph for cogview4 --- invokeai/backend/model_manager/taxonomy.py | 2 +- .../listeners/enqueueRequestedLinear.ts | 3 + .../util/graph/buildLinearBatchConfig.ts | 6 +- .../util/graph/generation/addImageToImage.ts | 18 +- .../nodes/util/graph/generation/addInpaint.ts | 23 ++- .../util/graph/generation/addNSFWChecker.ts | 15 +- .../util/graph/generation/addOutpaint.ts | 25 ++- .../util/graph/generation/addTextToImage.ts | 5 +- .../util/graph/generation/addWatermarker.ts | 15 +- .../graph/generation/buildCogView4Graph.ts | 187 ++++++++++++++++++ .../util/graph/generation/buildFLUXGraph.ts | 12 +- .../util/graph/generation/buildSD1Graph.ts | 12 +- .../util/graph/generation/buildSD3Graph.ts | 12 +- .../util/graph/generation/buildSDXLGraph.ts | 12 +- .../nodes/util/graph/graphBuilderUtils.ts | 10 + .../src/features/nodes/util/graph/types.ts | 33 ++++ .../frontend/web/src/services/api/schema.ts | 4 +- 17 files changed, 295 insertions(+), 99 deletions(-) create mode 100644 invokeai/frontend/web/src/features/nodes/util/graph/generation/buildCogView4Graph.ts create mode 100644 invokeai/frontend/web/src/features/nodes/util/graph/types.ts diff --git a/invokeai/backend/model_manager/taxonomy.py b/invokeai/backend/model_manager/taxonomy.py index 02420e8773..943c7f78bc 100644 --- a/invokeai/backend/model_manager/taxonomy.py +++ b/invokeai/backend/model_manager/taxonomy.py @@ -25,7 +25,7 @@ class BaseModelType(str, Enum): StableDiffusionXL = "sdxl" StableDiffusionXLRefiner = "sdxl-refiner" Flux = "flux" - CogView4 = " cogview4" + CogView4 = "cogview4" # Kandinsky2_1 = "kandinsky-2.1" diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/enqueueRequestedLinear.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/enqueueRequestedLinear.ts index 9de26aa07c..c357d87bb6 100644 --- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/enqueueRequestedLinear.ts +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/enqueueRequestedLinear.ts @@ -6,6 +6,7 @@ import { withResult, withResultAsync } from 'common/util/result'; import { parseify } from 'common/util/serialize'; import { $canvasManager } from 'features/controlLayers/store/ephemeral'; import { prepareLinearUIBatch } from 'features/nodes/util/graph/buildLinearBatchConfig'; +import { buildCogView4Graph } from 'features/nodes/util/graph/generation/buildCogView4Graph'; import { buildFLUXGraph } from 'features/nodes/util/graph/generation/buildFLUXGraph'; import { buildSD1Graph } from 'features/nodes/util/graph/generation/buildSD1Graph'; import { buildSD3Graph } from 'features/nodes/util/graph/generation/buildSD3Graph'; @@ -45,6 +46,8 @@ export const addEnqueueRequestedLinear = (startAppListening: AppStartListening) return await buildSD3Graph(state, manager); case `flux`: return await buildFLUXGraph(state, manager); + case 'cogview4': + return await buildCogView4Graph(state, manager); default: assert(false, `No graph builders for base ${base}`); } diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/buildLinearBatchConfig.ts b/invokeai/frontend/web/src/features/nodes/util/graph/buildLinearBatchConfig.ts index 94b0a1e5a7..46b1d6adc4 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/buildLinearBatchConfig.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/buildLinearBatchConfig.ts @@ -5,12 +5,14 @@ import { range } from 'lodash-es'; import type { components } from 'services/api/schema'; import type { Batch, EnqueueBatchArg, Invocation } from 'services/api/types'; +import type { ConditioningNodes, NoiseNodes } from './types'; + export const prepareLinearUIBatch = ( state: RootState, g: Graph, prepend: boolean, - noise: Invocation<'noise' | 'flux_denoise' | 'sd3_denoise'>, - posCond: Invocation<'compel' | 'sdxl_compel_prompt' | 'flux_text_encoder' | 'sd3_text_encoder'>, + noise: Invocation, + posCond: Invocation, origin: 'canvas' | 'workflows' | 'upscaling', destination: 'canvas' | 'gallery' ): EnqueueBatchArg => { diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/addImageToImage.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/addImageToImage.ts index 6813230f5c..07b15980f2 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/addImageToImage.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/addImageToImage.ts @@ -2,18 +2,22 @@ import type { CanvasManager } from 'features/controlLayers/konva/CanvasManager'; import { getPrefixedId } from 'features/controlLayers/konva/util'; import type { CanvasState, Dimensions } from 'features/controlLayers/store/types'; import type { Graph } from 'features/nodes/util/graph/generation/Graph'; +import type { + DenoiseLatentsNodes, + LatentToImageNodes, + MainModelLoaderNodes, + VaeSourceNodes, +} from 'features/nodes/util/graph/types'; import { isEqual } from 'lodash-es'; import type { Invocation } from 'services/api/types'; type AddImageToImageArg = { g: Graph; manager: CanvasManager; - l2i: Invocation<'l2i' | 'flux_vae_decode' | 'sd3_l2i'>; - i2lNodeType: 'i2l' | 'flux_vae_encode' | 'sd3_i2l'; - denoise: Invocation<'denoise_latents' | 'flux_denoise' | 'sd3_denoise'>; - vaeSource: Invocation< - 'main_model_loader' | 'sdxl_model_loader' | 'flux_model_loader' | 'seamless' | 'vae_loader' | 'sd3_model_loader' - >; + l2i: Invocation; + i2lNodeType: 'i2l' | 'flux_vae_encode' | 'sd3_i2l' | 'cogview4_i2l'; + denoise: Invocation; + vaeSource: Invocation; originalSize: Dimensions; scaledSize: Dimensions; bbox: CanvasState['bbox']; @@ -33,7 +37,7 @@ export const addImageToImage = async ({ bbox, denoising_start, fp32, -}: AddImageToImageArg): Promise> => { +}: AddImageToImageArg): Promise> => { denoise.denoising_start = denoising_start; const adapters = manager.compositor.getVisibleAdaptersOfType('raster_layer'); const { image_name } = await manager.compositor.getCompositeImageDTO(adapters, bbox.rect, { 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 deaf2814b8..500a75de78 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 @@ -6,6 +6,13 @@ import { selectParamsSlice } from 'features/controlLayers/store/paramsSlice'; import { selectCanvasSlice } from 'features/controlLayers/store/selectors'; import type { Dimensions } from 'features/controlLayers/store/types'; import type { Graph } from 'features/nodes/util/graph/generation/Graph'; +import { isMainModelWithoutUnet } from 'features/nodes/util/graph/graphBuilderUtils'; +import type { + DenoiseLatentsNodes, + LatentToImageNodes, + MainModelLoaderNodes, + VaeSourceNodes, +} from 'features/nodes/util/graph/types'; import { isEqual } from 'lodash-es'; import type { Invocation } from 'services/api/types'; @@ -13,13 +20,11 @@ type AddInpaintArg = { state: RootState; g: Graph; manager: CanvasManager; - l2i: Invocation<'l2i' | 'flux_vae_decode' | 'sd3_l2i'>; - i2lNodeType: 'i2l' | 'flux_vae_encode' | 'sd3_i2l'; - denoise: Invocation<'denoise_latents' | 'flux_denoise' | 'sd3_denoise'>; - vaeSource: Invocation< - 'main_model_loader' | 'sdxl_model_loader' | 'flux_model_loader' | 'seamless' | 'vae_loader' | 'sd3_model_loader' - >; - modelLoader: Invocation<'main_model_loader' | 'sdxl_model_loader' | 'flux_model_loader' | 'sd3_model_loader'>; + l2i: Invocation; + i2lNodeType: 'i2l' | 'flux_vae_encode' | 'sd3_i2l' | 'cogview4_i2l'; + denoise: Invocation; + vaeSource: Invocation; + modelLoader: Invocation; originalSize: Dimensions; scaledSize: Dimensions; denoising_start: number; @@ -119,7 +124,7 @@ export const addInpaint = async ({ g.addEdge(vaeSource, 'vae', i2l, 'vae'); g.addEdge(vaeSource, 'vae', createGradientMask, 'vae'); - if (modelLoader.type !== 'flux_model_loader' && modelLoader.type !== 'sd3_model_loader') { + if (!isMainModelWithoutUnet(modelLoader)) { g.addEdge(modelLoader, 'unet', createGradientMask, 'unet'); } g.addEdge(resizeImageToScaledSize, 'image', createGradientMask, 'image'); @@ -184,7 +189,7 @@ export const addInpaint = async ({ g.addEdge(i2l, 'latents', denoise, 'latents'); g.addEdge(vaeSource, 'vae', i2l, 'vae'); g.addEdge(vaeSource, 'vae', createGradientMask, 'vae'); - if (modelLoader.type !== 'flux_model_loader' && modelLoader.type !== 'sd3_model_loader') { + if (!isMainModelWithoutUnet(modelLoader)) { g.addEdge(modelLoader, 'unet', createGradientMask, 'unet'); } g.addEdge(createGradientMask, 'denoise_mask', denoise, 'denoise_mask'); diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/addNSFWChecker.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/addNSFWChecker.ts index 405e778352..58a00efda2 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/addNSFWChecker.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/addNSFWChecker.ts @@ -1,5 +1,6 @@ import { getPrefixedId } from 'features/controlLayers/konva/util'; import type { Graph } from 'features/nodes/util/graph/generation/Graph'; +import type { ImageOutputNodes } from 'features/nodes/util/graph/types'; import type { Invocation } from 'services/api/types'; /** @@ -8,19 +9,7 @@ import type { Invocation } from 'services/api/types'; * @param imageOutput The current image output node * @returns The nsfw checker node */ -export const addNSFWChecker = ( - g: Graph, - imageOutput: Invocation< - | 'l2i' - | 'img_nsfw' - | 'img_watermark' - | 'img_resize' - | 'invokeai_img_blend' - | 'apply_mask_to_image' - | 'flux_vae_decode' - | 'sd3_l2i' - > -): Invocation<'img_nsfw'> => { +export const addNSFWChecker = (g: Graph, imageOutput: Invocation): Invocation<'img_nsfw'> => { const nsfw = g.addNode({ type: 'img_nsfw', id: getPrefixedId('nsfw_checker'), 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 81d5b07f23..202d2d353b 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 @@ -6,7 +6,14 @@ import { selectParamsSlice } from 'features/controlLayers/store/paramsSlice'; import { selectCanvasSlice } from 'features/controlLayers/store/selectors'; import type { Dimensions } from 'features/controlLayers/store/types'; import type { Graph } from 'features/nodes/util/graph/generation/Graph'; -import { getInfill } from 'features/nodes/util/graph/graphBuilderUtils'; +import { getInfill, isMainModelWithoutUnet } from 'features/nodes/util/graph/graphBuilderUtils'; +import type { + DenoiseLatentsNodes, + ImageToLatentsNodes, + LatentToImageNodes, + MainModelLoaderNodes, + VaeSourceNodes, +} from 'features/nodes/util/graph/types'; import { isEqual } from 'lodash-es'; import type { Invocation } from 'services/api/types'; @@ -14,13 +21,11 @@ type AddOutpaintArg = { state: RootState; g: Graph; manager: CanvasManager; - l2i: Invocation<'l2i' | 'flux_vae_decode' | 'sd3_l2i'>; - i2lNodeType: 'i2l' | 'flux_vae_encode' | 'sd3_i2l'; - denoise: Invocation<'denoise_latents' | 'flux_denoise' | 'sd3_denoise'>; - vaeSource: Invocation< - 'main_model_loader' | 'sdxl_model_loader' | 'flux_model_loader' | 'seamless' | 'vae_loader' | 'sd3_model_loader' - >; - modelLoader: Invocation<'main_model_loader' | 'sdxl_model_loader' | 'flux_model_loader' | 'sd3_model_loader'>; + l2i: Invocation; + i2lNodeType: ImageToLatentsNodes; + denoise: Invocation; + vaeSource: Invocation; + modelLoader: Invocation; originalSize: Dimensions; scaledSize: Dimensions; denoising_start: number; @@ -116,7 +121,7 @@ export const addOutpaint = async ({ g.addEdge(infill, 'image', createGradientMask, 'image'); g.addEdge(resizeInputMaskToScaledSize, 'image', createGradientMask, 'mask'); g.addEdge(vaeSource, 'vae', createGradientMask, 'vae'); - if (modelLoader.type !== 'flux_model_loader' && modelLoader.type !== 'sd3_model_loader') { + if (!isMainModelWithoutUnet(modelLoader)) { g.addEdge(modelLoader, 'unet', createGradientMask, 'unet'); } @@ -216,7 +221,7 @@ export const addOutpaint = async ({ g.addEdge(i2l, 'latents', denoise, 'latents'); g.addEdge(vaeSource, 'vae', i2l, 'vae'); g.addEdge(vaeSource, 'vae', createGradientMask, 'vae'); - if (modelLoader.type !== 'flux_model_loader' && modelLoader.type !== 'sd3_model_loader') { + if (!isMainModelWithoutUnet(modelLoader)) { g.addEdge(modelLoader, 'unet', createGradientMask, 'unet'); } diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/addTextToImage.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/addTextToImage.ts index 2346f1eaab..5878aa9580 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/addTextToImage.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/addTextToImage.ts @@ -1,12 +1,13 @@ import { getPrefixedId } from 'features/controlLayers/konva/util'; import type { Dimensions } from 'features/controlLayers/store/types'; import type { Graph } from 'features/nodes/util/graph/generation/Graph'; +import type { LatentToImageNodes } from 'features/nodes/util/graph/types'; import { isEqual } from 'lodash-es'; import type { Invocation } from 'services/api/types'; type AddTextToImageArg = { g: Graph; - l2i: Invocation<'l2i' | 'flux_vae_decode' | 'sd3_l2i'>; + l2i: Invocation; originalSize: Dimensions; scaledSize: Dimensions; }; @@ -16,7 +17,7 @@ export const addTextToImage = ({ l2i, originalSize, scaledSize, -}: AddTextToImageArg): Invocation<'img_resize' | 'l2i' | 'flux_vae_decode' | 'sd3_l2i'> => { +}: AddTextToImageArg): Invocation<'img_resize' | 'l2i' | 'flux_vae_decode' | 'sd3_l2i' | 'cogview4_l2i'> => { if (!isEqual(scaledSize, originalSize)) { // We need to resize the output image back to the original size const resizeImageToOriginalSize = g.addNode({ diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/addWatermarker.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/addWatermarker.ts index a61b9a1cdb..d62e06eb9d 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/addWatermarker.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/addWatermarker.ts @@ -1,5 +1,6 @@ import { getPrefixedId } from 'features/controlLayers/konva/util'; import type { Graph } from 'features/nodes/util/graph/generation/Graph'; +import type { ImageOutputNodes } from 'features/nodes/util/graph/types'; import type { Invocation } from 'services/api/types'; /** @@ -8,19 +9,7 @@ import type { Invocation } from 'services/api/types'; * @param imageOutput The image output node * @returns The watermark node */ -export const addWatermarker = ( - g: Graph, - imageOutput: Invocation< - | 'l2i' - | 'img_nsfw' - | 'img_watermark' - | 'img_resize' - | 'invokeai_img_blend' - | 'apply_mask_to_image' - | 'flux_vae_decode' - | 'sd3_l2i' - > -): Invocation<'img_watermark'> => { +export const addWatermarker = (g: Graph, imageOutput: Invocation): Invocation<'img_watermark'> => { const watermark = g.addNode({ type: 'img_watermark', id: getPrefixedId('watermarker'), diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildCogView4Graph.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildCogView4Graph.ts new file mode 100644 index 0000000000..fb37ad66ba --- /dev/null +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildCogView4Graph.ts @@ -0,0 +1,187 @@ +import { logger } from 'app/logging/logger'; +import type { RootState } from 'app/store/store'; +import type { CanvasManager } from 'features/controlLayers/konva/CanvasManager'; +import { getPrefixedId } from 'features/controlLayers/konva/util'; +import { selectCanvasSettingsSlice } from 'features/controlLayers/store/canvasSettingsSlice'; +import { selectParamsSlice } from 'features/controlLayers/store/paramsSlice'; +import { selectCanvasMetadata, selectCanvasSlice } from 'features/controlLayers/store/selectors'; +import { fetchModelConfigWithTypeGuard } from 'features/metadata/util/modelFetchingHelpers'; +import { addImageToImage } from 'features/nodes/util/graph/generation/addImageToImage'; +import { addInpaint } from 'features/nodes/util/graph/generation/addInpaint'; +import { addNSFWChecker } from 'features/nodes/util/graph/generation/addNSFWChecker'; +import { addOutpaint } from 'features/nodes/util/graph/generation/addOutpaint'; +import { addTextToImage } from 'features/nodes/util/graph/generation/addTextToImage'; +import { addWatermarker } from 'features/nodes/util/graph/generation/addWatermarker'; +import { Graph } from 'features/nodes/util/graph/generation/Graph'; +import { + CANVAS_OUTPUT_PREFIX, + getBoardField, + getPresetModifiedPrompts, + getSizes, +} from 'features/nodes/util/graph/graphBuilderUtils'; +import type { ImageOutputNodes } from 'features/nodes/util/graph/types'; +import type { Invocation } from 'services/api/types'; +import { isNonRefinerMainModelConfig } from 'services/api/types'; +import type { Equals } from 'tsafe'; +import { assert } from 'tsafe'; + +const log = logger('system'); + +export const buildCogView4Graph = async ( + state: RootState, + manager: CanvasManager +): Promise<{ g: Graph; noise: Invocation<'cogview4_denoise'>; posCond: Invocation<'cogview4_text_encoder'> }> => { + const generationMode = await manager.compositor.getGenerationMode(); + log.debug({ generationMode }, 'Building CogView4 graph'); + + const params = selectParamsSlice(state); + const canvasSettings = selectCanvasSettingsSlice(state); + const canvas = selectCanvasSlice(state); + + const { bbox } = canvas; + + const { model, cfgScale: cfg_scale, seed, steps } = params; + + assert(model, 'No model found in state'); + + const { originalSize, scaledSize } = getSizes(bbox); + const { positivePrompt, negativePrompt } = getPresetModifiedPrompts(state); + + const g = new Graph(getPrefixedId('cogview4_graph')); + const modelLoader = g.addNode({ + type: 'cogview4_model_loader', + id: getPrefixedId('cogview4_model_loader'), + model, + }); + const posCond = g.addNode({ + type: 'cogview4_text_encoder', + id: getPrefixedId('pos_prompt'), + prompt: positivePrompt, + }); + + const negCond = g.addNode({ + type: 'cogview4_text_encoder', + id: getPrefixedId('neg_prompt'), + prompt: negativePrompt, + }); + + const denoise = g.addNode({ + type: 'cogview4_denoise', + id: getPrefixedId('denoise_latents'), + cfg_scale, + width: scaledSize.width, + height: scaledSize.height, + steps, + denoising_start: 0, + denoising_end: 1, + }); + const l2i = g.addNode({ + type: 'cogview4_l2i', + id: getPrefixedId('l2i'), + }); + + g.addEdge(modelLoader, 'transformer', denoise, 'transformer'); + g.addEdge(modelLoader, 'glm_encoder', posCond, 'glm_encoder'); + g.addEdge(modelLoader, 'glm_encoder', negCond, 'glm_encoder'); + g.addEdge(modelLoader, 'vae', l2i, 'vae'); + + g.addEdge(posCond, 'conditioning', denoise, 'positive_conditioning'); + g.addEdge(negCond, 'conditioning', denoise, 'negative_conditioning'); + + g.addEdge(denoise, 'latents', l2i, 'latents'); + + const modelConfig = await fetchModelConfigWithTypeGuard(model.key, isNonRefinerMainModelConfig); + assert(modelConfig.base === 'cogview4'); + + g.upsertMetadata({ + generation_mode: 'cogview4_txt2img', + cfg_scale, + width: originalSize.width, + height: originalSize.height, + positive_prompt: positivePrompt, + negative_prompt: negativePrompt, + model: Graph.getModelMetadataField(modelConfig), + seed, + steps, + }); + + const denoising_start = 1 - params.img2imgStrength; + + let canvasOutput: Invocation = l2i; + + if (generationMode === 'txt2img') { + canvasOutput = addTextToImage({ g, l2i, originalSize, scaledSize }); + } else if (generationMode === 'img2img') { + canvasOutput = await addImageToImage({ + g, + manager, + l2i, + i2lNodeType: 'cogview4_i2l', + denoise, + vaeSource: modelLoader, + originalSize, + scaledSize, + bbox, + denoising_start, + fp32: false, + }); + } else if (generationMode === 'inpaint') { + canvasOutput = await addInpaint({ + state, + g, + manager, + l2i, + i2lNodeType: 'cogview4_i2l', + denoise, + vaeSource: modelLoader, + modelLoader, + originalSize, + scaledSize, + denoising_start, + fp32: false, + }); + } else if (generationMode === 'outpaint') { + canvasOutput = await addOutpaint({ + state, + g, + manager, + l2i, + i2lNodeType: 'cogview4_i2l', + denoise, + vaeSource: modelLoader, + modelLoader, + originalSize, + scaledSize, + denoising_start, + fp32: false, + }); + } else { + assert>(false); + } + + if (state.system.shouldUseNSFWChecker) { + canvasOutput = addNSFWChecker(g, canvasOutput); + } + + if (state.system.shouldUseWatermarker) { + canvasOutput = addWatermarker(g, canvasOutput); + } + + // This image will be staged, should not be saved to the gallery or added to a board. + const is_intermediate = canvasSettings.sendToCanvas; + const board = canvasSettings.sendToCanvas ? undefined : getBoardField(state); + + if (!canvasSettings.sendToCanvas) { + g.upsertMetadata(selectCanvasMetadata(state)); + } + + g.updateNode(canvasOutput, { + id: getPrefixedId(CANVAS_OUTPUT_PREFIX), + is_intermediate, + use_cache: false, + board, + }); + + g.setMetadataReceivingNode(canvasOutput); + return { g, noise: denoise, posCond }; +}; diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.ts index 897b8e2c14..19c4b2e4fb 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildFLUXGraph.ts @@ -22,6 +22,7 @@ import { getPresetModifiedPrompts, getSizes, } from 'features/nodes/util/graph/graphBuilderUtils'; +import type { ImageOutputNodes } from 'features/nodes/util/graph/types'; import { t } from 'i18next'; import { selectMainModelConfig } from 'services/api/endpoints/models'; import type { Invocation } from 'services/api/types'; @@ -165,16 +166,7 @@ export const buildFLUXGraph = async ( denoising_start = 1 - img2imgStrength; } - let canvasOutput: Invocation< - | 'l2i' - | 'img_nsfw' - | 'img_watermark' - | 'img_resize' - | 'invokeai_img_blend' - | 'apply_mask_to_image' - | 'flux_vae_decode' - | 'sd3_l2i' - > = l2i; + let canvasOutput: Invocation = l2i; if (isFLUXFill) { canvasOutput = await addFLUXFill({ 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 7994146ed9..b34c4df7bf 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 @@ -23,6 +23,7 @@ import { getPresetModifiedPrompts, getSizes, } from 'features/nodes/util/graph/graphBuilderUtils'; +import type { ImageOutputNodes } from 'features/nodes/util/graph/types'; import { selectMainModelConfig } from 'services/api/endpoints/models'; import type { Invocation } from 'services/api/types'; import type { Equals } from 'tsafe'; @@ -167,16 +168,7 @@ export const buildSD1Graph = async ( const denoising_start = 1 - params.img2imgStrength; - let canvasOutput: Invocation< - | 'l2i' - | 'img_nsfw' - | 'img_watermark' - | 'img_resize' - | 'invokeai_img_blend' - | 'apply_mask_to_image' - | 'flux_vae_decode' - | 'sd3_l2i' - > = l2i; + let canvasOutput: Invocation = l2i; if (generationMode === 'txt2img') { canvasOutput = addTextToImage({ g, l2i, originalSize, scaledSize }); diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSD3Graph.ts b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSD3Graph.ts index 38288947d1..30f4f7235e 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSD3Graph.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSD3Graph.ts @@ -18,6 +18,7 @@ import { getPresetModifiedPrompts, getSizes, } from 'features/nodes/util/graph/graphBuilderUtils'; +import type { ImageOutputNodes } from 'features/nodes/util/graph/types'; import { selectMainModelConfig } from 'services/api/endpoints/models'; import type { Invocation } from 'services/api/types'; import type { Equals } from 'tsafe'; @@ -131,16 +132,7 @@ export const buildSD3Graph = async ( denoising_start = 1 - img2imgStrength; } - let canvasOutput: Invocation< - | 'l2i' - | 'img_nsfw' - | 'img_watermark' - | 'img_resize' - | 'invokeai_img_blend' - | 'apply_mask_to_image' - | 'flux_vae_decode' - | 'sd3_l2i' - > = l2i; + let canvasOutput: Invocation = l2i; if (generationMode === 'txt2img') { canvasOutput = addTextToImage({ g, l2i, originalSize, scaledSize }); 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 ff1fd45aed..1ec684ed02 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 @@ -23,6 +23,7 @@ import { getPresetModifiedPrompts, getSizes, } from 'features/nodes/util/graph/graphBuilderUtils'; +import type { ImageOutputNodes } from 'features/nodes/util/graph/types'; import { selectMainModelConfig } from 'services/api/endpoints/models'; import type { Invocation } from 'services/api/types'; import type { Equals } from 'tsafe'; @@ -173,16 +174,7 @@ export const buildSDXLGraph = async ( ? Math.min(refinerStart, 1 - params.img2imgStrength) : 1 - params.img2imgStrength; - let canvasOutput: Invocation< - | 'l2i' - | 'img_nsfw' - | 'img_watermark' - | 'img_resize' - | 'invokeai_img_blend' - | 'apply_mask_to_image' - | 'flux_vae_decode' - | 'sd3_l2i' - > = l2i; + let canvasOutput: Invocation = l2i; if (generationMode === 'txt2img') { canvasOutput = addTextToImage({ g, l2i, originalSize, scaledSize }); diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/graphBuilderUtils.ts b/invokeai/frontend/web/src/features/nodes/util/graph/graphBuilderUtils.ts index 3d1ad2c006..7686b99123 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/graphBuilderUtils.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/graphBuilderUtils.ts @@ -9,6 +9,8 @@ import { stylePresetsApi } from 'services/api/endpoints/stylePresets'; import type { Invocation } from 'services/api/types'; import { assert } from 'tsafe'; +import type { MainModelLoaderNodes } from './types'; + /** * Gets the board field, based on the autoAddBoardId setting. */ @@ -119,3 +121,11 @@ export const getInfill = ( }; export const CANVAS_OUTPUT_PREFIX = 'canvas_output'; + +export const isMainModelWithoutUnet = (modelLoader: Invocation) => { + return ( + modelLoader.type === 'flux_model_loader' || + modelLoader.type === 'sd3_model_loader' || + modelLoader.type === 'cogview4_model_loader' + ); +}; diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/types.ts b/invokeai/frontend/web/src/features/nodes/util/graph/types.ts new file mode 100644 index 0000000000..0d933d5653 --- /dev/null +++ b/invokeai/frontend/web/src/features/nodes/util/graph/types.ts @@ -0,0 +1,33 @@ +export type ImageOutputNodes = + | 'l2i' + | 'img_nsfw' + | 'img_watermark' + | 'img_resize' + | 'invokeai_img_blend' + | 'apply_mask_to_image' + | 'flux_vae_decode' + | 'sd3_l2i' + | 'cogview4_l2i'; + +export type LatentToImageNodes = 'l2i' | 'flux_vae_decode' | 'sd3_l2i' | 'cogview4_l2i'; + +export type ImageToLatentsNodes = 'i2l' | 'flux_vae_encode' | 'sd3_i2l' | 'cogview4_i2l'; + +export type DenoiseLatentsNodes = 'denoise_latents' | 'flux_denoise' | 'sd3_denoise' | 'cogview4_denoise'; + +export type MainModelLoaderNodes = + | 'main_model_loader' + | 'sdxl_model_loader' + | 'flux_model_loader' + | 'sd3_model_loader' + | 'cogview4_model_loader'; + +export type VaeSourceNodes = 'seamless' | 'vae_loader'; +export type NoiseNodes = 'noise' | 'flux_denoise' | 'sd3_denoise' | 'cogview4_denoise'; + +export type ConditioningNodes = + | 'compel' + | 'sdxl_compel_prompt' + | 'flux_text_encoder' + | 'sd3_text_encoder' + | 'cogview4_text_encoder'; diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index 3781516193..5213b4bb86 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -1970,7 +1970,7 @@ export type components = { * @description Base model type. * @enum {string} */ - BaseModelType: "any" | "sd-1" | "sd-2" | "sd-3" | "sdxl" | "sdxl-refiner" | "flux" | " cogview4"; + BaseModelType: "any" | "sd-1" | "sd-2" | "sd-3" | "sdxl" | "sdxl-refiner" | "flux" | "cogview4"; /** Batch */ Batch: { /** @@ -5119,7 +5119,7 @@ export type components = { * @description The generation mode that output this image * @default null */ - generation_mode?: ("txt2img" | "img2img" | "inpaint" | "outpaint" | "sdxl_txt2img" | "sdxl_img2img" | "sdxl_inpaint" | "sdxl_outpaint" | "flux_txt2img" | "flux_img2img" | "flux_inpaint" | "flux_outpaint" | "sd3_txt2img" | "sd3_img2img" | "sd3_inpaint" | "sd3_outpaint") | null; + generation_mode?: ("txt2img" | "img2img" | "inpaint" | "outpaint" | "sdxl_txt2img" | "sdxl_img2img" | "sdxl_inpaint" | "sdxl_outpaint" | "flux_txt2img" | "flux_img2img" | "flux_inpaint" | "flux_outpaint" | "sd3_txt2img" | "sd3_img2img" | "sd3_inpaint" | "sd3_outpaint" | "cogview4_txt2img" | "cogview4_img2img" | "cogview4_inpaint" | "cogview4_outpaint") | null; /** * Positive Prompt * @description The positive prompt parameter