feat (ui, generation): High Resolution Fix MVP in Text2Image Linear Flow (#4819)

* added HrfScale type with initial value

* working

* working

* working

* working

* working

* added addHrfToGraph

* continueing to implement this

* working on this

* comments

* working

* made hrf into its own collapse

* working on adding strength slider

* working

* working

* refactoring

* working

* change of this working: 0

* removed onnx support since apparently its not used

* working

* made scale integer

* trying out psycicpebbles idea

* working

* working on this

* working

* added toggle

* comments

* self review

* fixing things

* remove 'any' type

* fixing typing

* changed initial strength value to 3 (large values cause issues)

* set denoising start to be 1 - strength to resemble image to image

* set initial value

* added image to image

* pr1

* pr2

* updating to resolution finding

* working

* working

* working

* working

* working

* working

* working

* working

* working

* use memo

* connect rescale hw to noise

* working

* fixed min bug

* nit

* hides elements conditionally

* style

* feat(ui): add config for HRF, disable if feature disabled or ONNX model in use

* fix(ui): use `useCallback` for HRF toggle

---------

Co-authored-by: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
This commit is contained in:
Paul Curry
2023-10-14 03:34:41 -07:00
committed by GitHub
parent 40a568c060
commit 75663ec81e
18 changed files with 683 additions and 1 deletions

View File

@ -0,0 +1,247 @@
import { RootState } from 'app/store/store';
import { NonNullableGraph } from 'features/nodes/types/types';
import {
DenoiseLatentsInvocation,
ResizeLatentsInvocation,
NoiseInvocation,
LatentsToImageInvocation,
Edge,
} from 'services/api/types';
import {
LATENTS_TO_IMAGE,
DENOISE_LATENTS,
NOISE,
MAIN_MODEL_LOADER,
METADATA_ACCUMULATOR,
LATENTS_TO_IMAGE_HRF,
DENOISE_LATENTS_HRF,
RESCALE_LATENTS,
NOISE_HRF,
VAE_LOADER,
} from './constants';
import { logger } from 'app/logging/logger';
// Copy certain connections from previous DENOISE_LATENTS to new DENOISE_LATENTS_HRF.
function copyConnectionsToDenoiseLatentsHrf(graph: NonNullableGraph): void {
const destinationFields = [
'vae',
'control',
'ip_adapter',
'metadata',
'unet',
'positive_conditioning',
'negative_conditioning',
];
const newEdges: Edge[] = [];
// Loop through the existing edges connected to DENOISE_LATENTS
graph.edges.forEach((edge: Edge) => {
if (
edge.destination.node_id === DENOISE_LATENTS &&
destinationFields.includes(edge.destination.field)
) {
// Add a similar connection to DENOISE_LATENTS_HRF
newEdges.push({
source: {
node_id: edge.source.node_id,
field: edge.source.field,
},
destination: {
node_id: DENOISE_LATENTS_HRF,
field: edge.destination.field,
},
});
}
});
graph.edges = graph.edges.concat(newEdges);
}
// Adds the high-res fix feature to the given graph.
export const addHrfToGraph = (
state: RootState,
graph: NonNullableGraph
): void => {
// Double check hrf is enabled.
if (
!state.generation.hrfEnabled ||
state.config.disabledSDFeatures.includes('hrf') ||
state.generation.model?.model_type === 'onnx' // TODO: ONNX support
) {
return;
}
const log = logger('txt2img');
const { vae } = state.generation;
const isAutoVae = !vae;
const hrfWidth = state.generation.hrfWidth;
const hrfHeight = state.generation.hrfHeight;
// Pre-existing (original) graph nodes.
const originalDenoiseLatentsNode = graph.nodes[DENOISE_LATENTS] as
| DenoiseLatentsInvocation
| undefined;
const originalNoiseNode = graph.nodes[NOISE] as NoiseInvocation | undefined;
// Original latents to image should pick this up.
const originalLatentsToImageNode = graph.nodes[LATENTS_TO_IMAGE] as
| LatentsToImageInvocation
| undefined;
// Check if originalDenoiseLatentsNode is undefined and log an error
if (!originalDenoiseLatentsNode) {
log.error('originalDenoiseLatentsNode is undefined');
return;
}
// Check if originalNoiseNode is undefined and log an error
if (!originalNoiseNode) {
log.error('originalNoiseNode is undefined');
return;
}
// Check if originalLatentsToImageNode is undefined and log an error
if (!originalLatentsToImageNode) {
log.error('originalLatentsToImageNode is undefined');
return;
}
// Change height and width of original noise node to initial resolution.
if (originalNoiseNode) {
originalNoiseNode.width = hrfWidth;
originalNoiseNode.height = hrfHeight;
}
// Define new nodes.
// Denoise latents node to be run on upscaled latents.
const denoiseLatentsHrfNode: DenoiseLatentsInvocation = {
type: 'denoise_latents',
id: DENOISE_LATENTS_HRF,
is_intermediate: originalDenoiseLatentsNode?.is_intermediate,
cfg_scale: originalDenoiseLatentsNode?.cfg_scale,
scheduler: originalDenoiseLatentsNode?.scheduler,
steps: originalDenoiseLatentsNode?.steps,
denoising_start: 1 - state.generation.hrfStrength,
denoising_end: 1,
};
// New base resolution noise node.
const hrfNoiseNode: NoiseInvocation = {
type: 'noise',
id: NOISE_HRF,
seed: originalNoiseNode?.seed,
use_cpu: originalNoiseNode?.use_cpu,
is_intermediate: originalNoiseNode?.is_intermediate,
};
const rescaleLatentsNode: ResizeLatentsInvocation = {
id: RESCALE_LATENTS,
type: 'lresize',
width: state.generation.width,
height: state.generation.height,
};
// New node to convert latents to image.
const latentsToImageHrfNode: LatentsToImageInvocation | undefined =
originalLatentsToImageNode
? {
type: 'l2i',
id: LATENTS_TO_IMAGE_HRF,
fp32: originalLatentsToImageNode?.fp32,
is_intermediate: originalLatentsToImageNode?.is_intermediate,
}
: undefined;
// Add new nodes to graph.
graph.nodes[LATENTS_TO_IMAGE_HRF] =
latentsToImageHrfNode as LatentsToImageInvocation;
graph.nodes[DENOISE_LATENTS_HRF] =
denoiseLatentsHrfNode as DenoiseLatentsInvocation;
graph.nodes[NOISE_HRF] = hrfNoiseNode as NoiseInvocation;
graph.nodes[RESCALE_LATENTS] = rescaleLatentsNode as ResizeLatentsInvocation;
// Connect nodes.
graph.edges.push(
{
// Set up rescale latents.
source: {
node_id: DENOISE_LATENTS,
field: 'latents',
},
destination: {
node_id: RESCALE_LATENTS,
field: 'latents',
},
},
// Set up new noise node
{
source: {
node_id: RESCALE_LATENTS,
field: 'height',
},
destination: {
node_id: NOISE_HRF,
field: 'height',
},
},
{
source: {
node_id: RESCALE_LATENTS,
field: 'width',
},
destination: {
node_id: NOISE_HRF,
field: 'width',
},
},
// Set up new denoise node.
{
source: {
node_id: RESCALE_LATENTS,
field: 'latents',
},
destination: {
node_id: DENOISE_LATENTS_HRF,
field: 'latents',
},
},
{
source: {
node_id: NOISE_HRF,
field: 'noise',
},
destination: {
node_id: DENOISE_LATENTS_HRF,
field: 'noise',
},
},
// Set up new latents to image node.
{
source: {
node_id: DENOISE_LATENTS_HRF,
field: 'latents',
},
destination: {
node_id: LATENTS_TO_IMAGE_HRF,
field: 'latents',
},
},
{
source: {
node_id: METADATA_ACCUMULATOR,
field: 'metadata',
},
destination: {
node_id: LATENTS_TO_IMAGE_HRF,
field: 'metadata',
},
},
{
source: {
node_id: isAutoVae ? MAIN_MODEL_LOADER : VAE_LOADER,
field: 'vae',
},
destination: {
node_id: LATENTS_TO_IMAGE_HRF,
field: 'vae',
},
}
);
copyConnectionsToDenoiseLatentsHrf(graph);
};

View File

@ -2,6 +2,7 @@ import { NonNullableGraph } from 'features/nodes/types/types';
import {
CANVAS_OUTPUT,
LATENTS_TO_IMAGE,
LATENTS_TO_IMAGE_HRF,
METADATA_ACCUMULATOR,
NSFW_CHECKER,
SAVE_IMAGE,
@ -82,6 +83,14 @@ export const addSaveImageNode = (
},
destination,
});
} else if (LATENTS_TO_IMAGE_HRF in graph.nodes) {
graph.edges.push({
source: {
node_id: LATENTS_TO_IMAGE_HRF,
field: 'image',
},
destination,
});
} else if (LATENTS_TO_IMAGE in graph.nodes) {
graph.edges.push({
source: {

View File

@ -6,6 +6,7 @@ import {
ONNXTextToLatentsInvocation,
} from 'services/api/types';
import { addControlNetToLinearGraph } from './addControlNetToLinearGraph';
import { addHrfToGraph } from './addHrfToGraph';
import { addIPAdapterToLinearGraph } from './addIPAdapterToLinearGraph';
import { addLoRAsToGraph } from './addLoRAsToGraph';
import { addNSFWCheckerToGraph } from './addNSFWCheckerToGraph';
@ -47,6 +48,10 @@ export const buildLinearTextToImageGraph = (
seamlessXAxis,
seamlessYAxis,
seed,
hrfWidth,
hrfHeight,
hrfStrength,
hrfEnabled: hrfEnabled,
} = state.generation;
const use_cpu = shouldUseCpuNoise;
@ -254,6 +259,9 @@ export const buildLinearTextToImageGraph = (
ipAdapters: [], // populated in addIPAdapterToLinearGraph
t2iAdapters: [], // populated in addT2IAdapterToLinearGraph
clip_skip: clipSkip,
hrf_width: hrfEnabled ? hrfWidth : undefined,
hrf_height: hrfEnabled ? hrfHeight : undefined,
hrf_strength: hrfEnabled ? hrfStrength : undefined,
};
graph.edges.push({
@ -287,6 +295,12 @@ export const buildLinearTextToImageGraph = (
addT2IAdaptersToLinearGraph(state, graph, DENOISE_LATENTS);
// High resolution fix.
// NOTE: Not supported for onnx models.
if (state.generation.hrfEnabled && !isUsingOnnxModel) {
addHrfToGraph(state, graph);
}
// NSFW & watermark - must be last thing added to graph
if (state.system.shouldUseNSFWChecker) {
// must add before watermarker!

View File

@ -2,11 +2,14 @@
export const POSITIVE_CONDITIONING = 'positive_conditioning';
export const NEGATIVE_CONDITIONING = 'negative_conditioning';
export const DENOISE_LATENTS = 'denoise_latents';
export const DENOISE_LATENTS_HRF = 'denoise_latents_hrf';
export const LATENTS_TO_IMAGE = 'latents_to_image';
export const LATENTS_TO_IMAGE_HRF = 'latents_to_image_hrf';
export const SAVE_IMAGE = 'save_image';
export const NSFW_CHECKER = 'nsfw_checker';
export const WATERMARKER = 'invisible_watermark';
export const NOISE = 'noise';
export const NOISE_HRF = 'noise_hrf';
export const RANDOM_INT = 'rand_int';
export const RANGE_OF_SIZE = 'range_of_size';
export const ITERATE = 'iterate';
@ -18,6 +21,7 @@ export const CLIP_SKIP = 'clip_skip';
export const IMAGE_TO_LATENTS = 'image_to_latents';
export const LATENTS_TO_LATENTS = 'latents_to_latents';
export const RESIZE = 'resize_image';
export const RESCALE_LATENTS = 'rescale_latents';
export const IMG2IMG_RESIZE = 'img2img_resize';
export const CANVAS_OUTPUT = 'canvas_output';
export const INPAINT_IMAGE = 'inpaint_image';