Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 19 additions & 1 deletion invokeai/app/invocations/image_to_latents.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from contextlib import nullcontext
from functools import singledispatchmethod
from typing import Literal

import einops
import torch
Expand Down Expand Up @@ -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."""
Expand All @@ -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(
Expand Down Expand Up @@ -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")

Expand Down
5 changes: 5 additions & 0 deletions invokeai/frontend/web/public/locales/en.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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": [
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ export type Feature =
| 'scaleBeforeProcessing'
| 'seamlessTilingXAxis'
| 'seamlessTilingYAxis'
| 'colorCompensation'
| 'upscaleModel'
| 'scale'
| 'creativity'
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
},
Expand Down Expand Up @@ -436,6 +439,7 @@ export const {
clipGEmbedModelSelected,
setClipSkip,
shouldUseCpuNoiseChanged,
setColorCompensation,
positivePromptChanged,
positivePromptAddedToHistory,
promptRemovedFromHistory,
Expand Down Expand Up @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -645,6 +646,7 @@ export const getInitialParamsState = (): ParamsState => ({
seamlessYAxis: false,
clipSkip: 0,
shouldUseCpuNoise: true,
colorCompensation: false,
positivePrompt: '',
positivePromptHistory: [],
negativePrompt: null,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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'));
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down
Original file line number Diff line number Diff line change
@@ -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);
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -97,6 +98,9 @@ export const AdvancedSettingsAccordion = memo(() => {
<ParamSeamlessYAxis />
</FormControlGroup>
</Flex>
<FormControlGroup formLabelProps={formLabelProps}>
<ParamColorCompensation />
</FormControlGroup>
</>
)}
{isFLUX && (
Expand Down
7 changes: 7 additions & 0 deletions invokeai/frontend/web/src/services/api/schema.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down