From 490c77d2a78fd0d22ed996f006b86a8e443d5993 Mon Sep 17 00:00:00 2001 From: dunkeroni Date: Sat, 25 Oct 2025 19:11:52 -0400 Subject: [PATCH 1/6] feat(nodes/UI): add SDXL color compensation option --- invokeai/app/invocations/image_to_latents.py | 20 +++++++++++- invokeai/frontend/web/public/locales/en.json | 5 +++ .../InformationalPopover/constants.ts | 1 + .../controlLayers/store/paramsSlice.ts | 5 +++ .../src/features/controlLayers/store/types.ts | 2 ++ .../util/graph/generation/buildSDXLGraph.ts | 5 +++ .../VAEModel/ParamColorCompensation.tsx | 32 +++++++++++++++++++ .../AdvancedSettingsAccordion.tsx | 4 +++ .../frontend/web/src/services/api/schema.ts | 7 ++++ 9 files changed, 80 insertions(+), 1 deletion(-) create mode 100644 invokeai/frontend/web/src/features/parameters/components/VAEModel/ParamColorCompensation.tsx diff --git a/invokeai/app/invocations/image_to_latents.py b/invokeai/app/invocations/image_to_latents.py index fde70a34fde..cc1ed5d1f20 100644 --- a/invokeai/app/invocations/image_to_latents.py +++ b/invokeai/app/invocations/image_to_latents.py @@ -1,5 +1,6 @@ from contextlib import nullcontext from functools import singledispatchmethod +from typing import Literal import einops import torch @@ -29,13 +30,21 @@ from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_sd15_sdxl +""" +SDXL VAE color compensation values determined experimentally to reduce color drift by a factor of ~1/5. +If more precise values are found in the future (e.g. individual color channels), they can be updated. +SD1.5, TAESD, TAESDXL VAEs distort in less predictable ways, so no compensation is offered at this time. +""" +COMPENSATION_OPTIONS = Literal["None", "SDXL"] +COLOR_COMPENSATION_MAP = {"None": [1, 0], "SDXL": [1.02, -0.002]} + @invocation( "i2l", title="Image to Latents - SD1.5, SDXL", tags=["latents", "image", "vae", "i2l"], category="latents", - version="1.1.1", + version="1.2.0", ) class ImageToLatentsInvocation(BaseInvocation): """Encodes an image into latents.""" @@ -52,6 +61,10 @@ class ImageToLatentsInvocation(BaseInvocation): # offer a way to directly set None values. tile_size: int = InputField(default=0, multiple_of=8, description=FieldDescriptions.vae_tile_size) fp32: bool = InputField(default=False, description=FieldDescriptions.fp32) + color_compensation: COMPENSATION_OPTIONS = InputField( + default="None", + description="Apply VAE scaling compensation when encoding images (reduces color drift).", + ) @classmethod def vae_encode( @@ -130,6 +143,11 @@ def invoke(self, context: InvocationContext) -> LatentsOutput: assert isinstance(vae_info.model, (AutoencoderKL, AutoencoderTiny)) image_tensor = image_resized_to_grid_as_tensor(image.convert("RGB")) + + if self.color_compensation != "None": + scale, bias = COLOR_COMPENSATION_MAP[self.color_compensation] + image_tensor = image_tensor * scale + bias + if image_tensor.dim() == 3: image_tensor = einops.rearrange(image_tensor, "c h w -> 1 c h w") diff --git a/invokeai/frontend/web/public/locales/en.json b/invokeai/frontend/web/public/locales/en.json index 8a6bd7b337e..dd19eb5473f 100644 --- a/invokeai/frontend/web/public/locales/en.json +++ b/invokeai/frontend/web/public/locales/en.json @@ -1317,6 +1317,7 @@ "scheduler": "Scheduler", "seamlessXAxis": "Seamless X Axis", "seamlessYAxis": "Seamless Y Axis", + "colorCompensation": "Color Compensation", "seed": "Seed", "imageActions": "Image Actions", "sendToCanvas": "Send To Canvas", @@ -1860,6 +1861,10 @@ "heading": "Seamless Tiling Y Axis", "paragraphs": ["Seamlessly tile an image along the vertical axis."] }, + "colorCompensation": { + "heading": "Color Compensation", + "paragraphs": ["Adjust the input image to reduce color shifts during inpainting or img2img (SDXL Only)."] + }, "upscaleModel": { "heading": "Upscale Model", "paragraphs": [ diff --git a/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts b/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts index 6db4dcbd682..89b21726675 100644 --- a/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts +++ b/invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts @@ -62,6 +62,7 @@ export type Feature = | 'scaleBeforeProcessing' | 'seamlessTilingXAxis' | 'seamlessTilingYAxis' + | 'colorCompensation' | 'upscaleModel' | 'scale' | 'creativity' diff --git a/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts b/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts index 9dd85b1bc20..ecdd70f3cad 100644 --- a/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts +++ b/invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts @@ -170,6 +170,9 @@ const slice = createSlice({ shouldUseCpuNoiseChanged: (state, action: PayloadAction) => { state.shouldUseCpuNoise = action.payload; }, + setColorCompensation: (state, action: PayloadAction) => { + state.colorCompensation = action.payload; + }, positivePromptChanged: (state, action: PayloadAction) => { state.positivePrompt = action.payload; }, @@ -436,6 +439,7 @@ export const { clipGEmbedModelSelected, setClipSkip, shouldUseCpuNoiseChanged, + setColorCompensation, positivePromptChanged, positivePromptAddedToHistory, promptRemovedFromHistory, @@ -557,6 +561,7 @@ export const selectShouldRandomizeSeed = createParamsSelector((params) => params export const selectVAEPrecision = createParamsSelector((params) => params.vaePrecision); export const selectIterations = createParamsSelector((params) => params.iterations); export const selectShouldUseCPUNoise = createParamsSelector((params) => params.shouldUseCpuNoise); +export const selectColorCompensation = createParamsSelector((params) => params.colorCompensation); export const selectUpscaleScheduler = createParamsSelector((params) => params.upscaleScheduler); export const selectUpscaleCfgScale = createParamsSelector((params) => params.upscaleCfgScale); diff --git a/invokeai/frontend/web/src/features/controlLayers/store/types.ts b/invokeai/frontend/web/src/features/controlLayers/store/types.ts index 87c173d7cca..4d6860f7278 100644 --- a/invokeai/frontend/web/src/features/controlLayers/store/types.ts +++ b/invokeai/frontend/web/src/features/controlLayers/store/types.ts @@ -596,6 +596,7 @@ export const zParamsState = z.object({ seamlessYAxis: z.boolean(), clipSkip: z.number(), shouldUseCpuNoise: z.boolean(), + colorCompensation: z.boolean(), positivePrompt: zParameterPositivePrompt, positivePromptHistory: zPositivePromptHistory, negativePrompt: zParameterNegativePrompt, @@ -645,6 +646,7 @@ export const getInitialParamsState = (): ParamsState => ({ seamlessYAxis: false, clipSkip: 0, shouldUseCpuNoise: true, + colorCompensation: false, positivePrompt: '', positivePromptHistory: [], negativePrompt: null, 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 1a4b3f9e963..9d65076a70d 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 @@ -45,12 +45,14 @@ export const buildSDXLGraph = async (arg: GraphBuilderArg): Promise { + const { t } = useTranslation(); + const colorCompensation = useAppSelector(selectColorCompensation); + + const dispatch = useAppDispatch(); + + const handleChange = useCallback( + (e: ChangeEvent) => { + dispatch(setColorCompensation(e.target.checked)); + }, + [dispatch] + ); + + return ( + + + {t('parameters.colorCompensation')} + + + + ); +}; + +export default memo(ParamColorCompensation); diff --git a/invokeai/frontend/web/src/features/settingsAccordions/components/AdvancedSettingsAccordion/AdvancedSettingsAccordion.tsx b/invokeai/frontend/web/src/features/settingsAccordions/components/AdvancedSettingsAccordion/AdvancedSettingsAccordion.tsx index 37475ad6aa8..2212cd153cd 100644 --- a/invokeai/frontend/web/src/features/settingsAccordions/components/AdvancedSettingsAccordion/AdvancedSettingsAccordion.tsx +++ b/invokeai/frontend/web/src/features/settingsAccordions/components/AdvancedSettingsAccordion/AdvancedSettingsAccordion.tsx @@ -12,6 +12,7 @@ import ParamClipSkip from 'features/parameters/components/Advanced/ParamClipSkip import ParamT5EncoderModelSelect from 'features/parameters/components/Advanced/ParamT5EncoderModelSelect'; import ParamSeamlessXAxis from 'features/parameters/components/Seamless/ParamSeamlessXAxis'; import ParamSeamlessYAxis from 'features/parameters/components/Seamless/ParamSeamlessYAxis'; +import ParamColorCompensation from 'features/parameters/components/VAEModel/ParamColorCompensation'; import ParamFLUXVAEModelSelect from 'features/parameters/components/VAEModel/ParamFLUXVAEModelSelect'; import ParamVAEModelSelect from 'features/parameters/components/VAEModel/ParamVAEModelSelect'; import ParamVAEPrecision from 'features/parameters/components/VAEModel/ParamVAEPrecision'; @@ -97,6 +98,9 @@ export const AdvancedSettingsAccordion = memo(() => { + + + )} {isFLUX && ( diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index 40214ffa554..d8b47dd58d2 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -11676,6 +11676,13 @@ export type components = { * @default false */ fp32?: boolean; + /** + * Color Compensation + * @description Apply VAE scaling compensation when encoding images (reduces color drift). + * @default None + * @enum {string} + */ + color_compensation?: "None" | "SDXL"; /** * type * @default i2l From 8a79b97d6e72865e71a48ff6ab754ed8b6a0e846 Mon Sep 17 00:00:00 2001 From: dunkeroni Date: Sat, 25 Oct 2025 19:48:54 -0400 Subject: [PATCH 2/6] adjust value --- invokeai/app/invocations/image_to_latents.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/invokeai/app/invocations/image_to_latents.py b/invokeai/app/invocations/image_to_latents.py index cc1ed5d1f20..148286fba84 100644 --- a/invokeai/app/invocations/image_to_latents.py +++ b/invokeai/app/invocations/image_to_latents.py @@ -31,12 +31,12 @@ from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_sd15_sdxl """ -SDXL VAE color compensation values determined experimentally to reduce color drift by a factor of ~1/5. -If more precise values are found in the future (e.g. individual color channels), they can be updated. +SDXL VAE color compensation values determined experimentally to reduce color drift. +If more reliable values are found in the future (e.g. individual color channels), they can be updated. SD1.5, TAESD, TAESDXL VAEs distort in less predictable ways, so no compensation is offered at this time. """ COMPENSATION_OPTIONS = Literal["None", "SDXL"] -COLOR_COMPENSATION_MAP = {"None": [1, 0], "SDXL": [1.02, -0.002]} +COLOR_COMPENSATION_MAP = {"None": [1, 0], "SDXL": [1.015, -0.002]} @invocation( From d991da710d4261a53e61a9f7dde51222c9da871d Mon Sep 17 00:00:00 2001 From: dunkeroni Date: Sat, 8 Nov 2025 18:52:56 -0500 Subject: [PATCH 3/6] Better warnings on wrong VAE base model --- invokeai/app/invocations/image_to_latents.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/invokeai/app/invocations/image_to_latents.py b/invokeai/app/invocations/image_to_latents.py index 148286fba84..600202defe6 100644 --- a/invokeai/app/invocations/image_to_latents.py +++ b/invokeai/app/invocations/image_to_latents.py @@ -75,7 +75,7 @@ def vae_encode( image_tensor: torch.Tensor, tile_size: int = 0, ) -> torch.Tensor: - assert isinstance(vae_info.model, (AutoencoderKL, AutoencoderTiny)) + assert isinstance(vae_info.model, (AutoencoderKL, AutoencoderTiny)), "VAE must be of type SD-1.5 or SDXL" estimated_working_memory = estimate_vae_working_memory_sd15_sdxl( operation="encode", image_tensor=image_tensor, @@ -84,7 +84,7 @@ def vae_encode( fp32=upcast, ) with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): - assert isinstance(vae, (AutoencoderKL, AutoencoderTiny)) + assert isinstance(vae, (AutoencoderKL, AutoencoderTiny)), "VAE must be of type SD-1.5 or SDXL" orig_dtype = vae.dtype if upcast: vae.to(dtype=torch.float32) @@ -140,7 +140,7 @@ def invoke(self, context: InvocationContext) -> LatentsOutput: image = context.images.get_pil(self.image.image_name) vae_info = context.models.load(self.vae.vae) - assert isinstance(vae_info.model, (AutoencoderKL, AutoencoderTiny)) + assert isinstance(vae_info.model, (AutoencoderKL, AutoencoderTiny)), "VAE must be of type SD-1.5 or SDXL" image_tensor = image_resized_to_grid_as_tensor(image.convert("RGB")) From 04de466a2cb996dc15b33e3caf6a221c2d2bc126 Mon Sep 17 00:00:00 2001 From: dunkeroni Date: Sat, 8 Nov 2025 18:57:02 -0500 Subject: [PATCH 4/6] Restrict XL compensation to XL models Co-authored-by: Lincoln Stein --- invokeai/app/invocations/image_to_latents.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/invokeai/app/invocations/image_to_latents.py b/invokeai/app/invocations/image_to_latents.py index 600202defe6..44498e1a3d7 100644 --- a/invokeai/app/invocations/image_to_latents.py +++ b/invokeai/app/invocations/image_to_latents.py @@ -144,7 +144,7 @@ def invoke(self, context: InvocationContext) -> LatentsOutput: image_tensor = image_resized_to_grid_as_tensor(image.convert("RGB")) - if self.color_compensation != "None": + if self.color_compensation != "None" and vae_info.config.base == BaseModelType.StableDiffusionXL: scale, bias = COLOR_COMPENSATION_MAP[self.color_compensation] image_tensor = image_tensor * scale + bias From be48c12ff1f0db88592583eeefcee147b2b7bdc4 Mon Sep 17 00:00:00 2001 From: dunkeroni Date: Sat, 8 Nov 2025 18:59:43 -0500 Subject: [PATCH 5/6] fix: BaseModelType missing import --- invokeai/app/invocations/image_to_latents.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/invokeai/app/invocations/image_to_latents.py b/invokeai/app/invocations/image_to_latents.py index 44498e1a3d7..bc08eebcc3b 100644 --- a/invokeai/app/invocations/image_to_latents.py +++ b/invokeai/app/invocations/image_to_latents.py @@ -21,7 +21,7 @@ Input, InputField, ) -from invokeai.app.invocations.model import VAEField +from invokeai.app.invocations.model import VAEField, BaseModelType from invokeai.app.invocations.primitives import LatentsOutput from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.model_manager.load.load_base import LoadedModel From 5a44414ca6126605fe5b73def912e9581c83fe60 Mon Sep 17 00:00:00 2001 From: dunkeroni Date: Sat, 8 Nov 2025 19:02:04 -0500 Subject: [PATCH 6/6] (chore): appease the ruff --- invokeai/app/invocations/image_to_latents.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/invokeai/app/invocations/image_to_latents.py b/invokeai/app/invocations/image_to_latents.py index bc08eebcc3b..8dc5ceba0b0 100644 --- a/invokeai/app/invocations/image_to_latents.py +++ b/invokeai/app/invocations/image_to_latents.py @@ -21,7 +21,7 @@ Input, InputField, ) -from invokeai.app.invocations.model import VAEField, BaseModelType +from invokeai.app.invocations.model import BaseModelType, VAEField from invokeai.app.invocations.primitives import LatentsOutput from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.model_manager.load.load_base import LoadedModel