diff --git a/invokeai/app/invocations/image_to_latents.py b/invokeai/app/invocations/image_to_latents.py index fde70a34fde..148286fba84 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. +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( @@ -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