From 490c77d2a78fd0d22ed996f006b86a8e443d5993 Mon Sep 17 00:00:00 2001 From: dunkeroni Date: Sat, 25 Oct 2025 19:11:52 -0400 Subject: [PATCH 1/2] 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/2] 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(