Feat: SDXL Color Compensation (#8637)

* feat(nodes/UI): add SDXL color compensation option

* adjust value

* Better warnings on wrong VAE base model

* Restrict XL compensation to XL models

Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com>

* fix: BaseModelType missing import

* (chore): appease the ruff

---------

Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com>

Author: dunkeroni

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
@@ -20,7 +21,7 @@
     Input,
     InputField,
 )
-from invokeai.app.invocations.model import VAEField
+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
@@ -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.
+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.015, -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(
@@ -62,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,
@@ -71,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)
@@ -127,9 +140,14 @@ 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"))
+
+        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
+
         if image_tensor.dim() == 3:
             image_tensor = einops.rearrange(image_tensor, "c h w -> 1 c h w")
 

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": [

invokeai/frontend/web/src/common/components/InformationalPopover/constants.ts

@@ -62,6 +62,7 @@ export type Feature =
   | 'scaleBeforeProcessing'
   | 'seamlessTilingXAxis'
   | 'seamlessTilingYAxis'
+  | 'colorCompensation'
   | 'upscaleModel'
   | 'scale'
   | 'creativity'

invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts

@@ -170,6 +170,9 @@ const slice = createSlice({
     shouldUseCpuNoiseChanged: (state, action: PayloadAction<boolean>) => {
       state.shouldUseCpuNoise = action.payload;
     },
+    setColorCompensation: (state, action: PayloadAction<boolean>) => {
+      state.colorCompensation = action.payload;
+    },
     positivePromptChanged: (state, action: PayloadAction<ParameterPositivePrompt>) => {
       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);

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,

invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSDXLGraph.ts

@@ -45,12 +45,14 @@ export const buildSDXLGraph = async (arg: GraphBuilderArg): Promise<GraphBuilder
     scheduler,
     steps,
     shouldUseCpuNoise,
+    colorCompensation,
     vaePrecision,
     vae,
     refinerModel,
   } = params;
 
   const fp32 = vaePrecision === 'fp32';
+  const compensation = colorCompensation ? 'SDXL' : 'None';
   const prompts = selectPresetModifiedPrompts(state);
 
   const g = new Graph(getPrefixedId('sdxl_graph'));
@@ -178,6 +180,7 @@ export const buildSDXLGraph = async (arg: GraphBuilderArg): Promise<GraphBuilder
       type: 'i2l',
       id: getPrefixedId('i2l'),
       fp32,
+      color_compensation: compensation,
     });
     canvasOutput = await addImageToImage({
       g,
@@ -196,6 +199,7 @@ export const buildSDXLGraph = async (arg: GraphBuilderArg): Promise<GraphBuilder
       type: 'i2l',
       id: getPrefixedId('i2l'),
       fp32,
+      color_compensation: compensation,
     });
     canvasOutput = await addInpaint({
       g,
@@ -216,6 +220,7 @@ export const buildSDXLGraph = async (arg: GraphBuilderArg): Promise<GraphBuilder
       type: 'i2l',
       id: getPrefixedId('i2l'),
       fp32,
+      color_compensation: compensation,
     });
     canvasOutput = await addOutpaint({
       g,

invokeai/frontend/web/src/features/parameters/components/VAEModel/ParamColorCompensation.tsx

@@ -0,0 +1,32 @@
+import { FormControl, FormLabel, Switch } from '@invoke-ai/ui-library';
+import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { InformationalPopover } from 'common/components/InformationalPopover/InformationalPopover';
+import { selectColorCompensation, setColorCompensation } from 'features/controlLayers/store/paramsSlice';
+import type { ChangeEvent } from 'react';
+import { memo, useCallback } from 'react';
+import { useTranslation } from 'react-i18next';
+
+const ParamColorCompensation = () => {
+  const { t } = useTranslation();
+  const colorCompensation = useAppSelector(selectColorCompensation);
+
+  const dispatch = useAppDispatch();
+
+  const handleChange = useCallback(
+    (e: ChangeEvent<HTMLInputElement>) => {
+      dispatch(setColorCompensation(e.target.checked));
+    },
+    [dispatch]
+  );
+
+  return (
+    <FormControl>
+      <InformationalPopover feature="colorCompensation">
+        <FormLabel>{t('parameters.colorCompensation')}</FormLabel>
+      </InformationalPopover>
+      <Switch isChecked={colorCompensation} onChange={handleChange} />
+    </FormControl>
+  );
+};
+
+export default memo(ParamColorCompensation);

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(() => {
                 <ParamSeamlessYAxis />
               </FormControlGroup>
             </Flex>
+            <FormControlGroup formLabelProps={formLabelProps}>
+              <ParamColorCompensation />
+            </FormControlGroup>
           </>
         )}
         {isFLUX && (

invokeai/frontend/web/src/services/api/schema.ts

@@ -11685,6 +11685,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