feat(ui): refactor parameter recall

- use zod to validate parameters before recalling
- update recall params hook to handle all validation and UI feedback
This commit is contained in:
psychedelicious
2023-05-31 21:05:31 +10:00
parent 062b2cf46f
commit 6571e4c2fd
15 changed files with 667 additions and 764 deletions

View File

@ -1,151 +0,0 @@
import { useAppDispatch } from 'app/store/storeHooks';
import { isFinite, isString } from 'lodash-es';
import { useCallback } from 'react';
import { useTranslation } from 'react-i18next';
import useSetBothPrompts from './usePrompt';
import { allParametersSet, setSeed } from '../store/generationSlice';
import { isImageField } from 'services/types/guards';
import { NUMPY_RAND_MAX } from 'app/constants';
import { initialImageSelected } from '../store/actions';
import { setActiveTab } from 'features/ui/store/uiSlice';
import { useAppToaster } from 'app/components/Toaster';
import { ImageDTO } from 'services/api';
export const useParameters = () => {
const dispatch = useAppDispatch();
const toaster = useAppToaster();
const { t } = useTranslation();
const setBothPrompts = useSetBothPrompts();
/**
* Sets prompt with toast
*/
const recallPrompt = useCallback(
(prompt: unknown, negativePrompt?: unknown) => {
if (!isString(prompt) || !isString(negativePrompt)) {
toaster({
title: t('toast.promptNotSet'),
description: t('toast.promptNotSetDesc'),
status: 'warning',
duration: 2500,
isClosable: true,
});
return;
}
setBothPrompts(prompt, negativePrompt);
toaster({
title: t('toast.promptSet'),
status: 'info',
duration: 2500,
isClosable: true,
});
},
[t, toaster, setBothPrompts]
);
/**
* Sets seed with toast
*/
const recallSeed = useCallback(
(seed: unknown) => {
const s = Number(seed);
if (!isFinite(s) || (isFinite(s) && !(s >= 0 && s <= NUMPY_RAND_MAX))) {
toaster({
title: t('toast.seedNotSet'),
description: t('toast.seedNotSetDesc'),
status: 'warning',
duration: 2500,
isClosable: true,
});
return;
}
dispatch(setSeed(s));
toaster({
title: t('toast.seedSet'),
status: 'info',
duration: 2500,
isClosable: true,
});
},
[t, toaster, dispatch]
);
/**
* Sets initial image with toast
*/
const recallInitialImage = useCallback(
async (image: unknown) => {
if (!isImageField(image)) {
toaster({
title: t('toast.initialImageNotSet'),
description: t('toast.initialImageNotSetDesc'),
status: 'warning',
duration: 2500,
isClosable: true,
});
return;
}
dispatch(initialImageSelected(image.image_name));
toaster({
title: t('toast.initialImageSet'),
status: 'info',
duration: 2500,
isClosable: true,
});
},
[t, toaster, dispatch]
);
/**
* Sets image as initial image with toast
*/
const sendToImageToImage = useCallback(
(image: ImageDTO) => {
dispatch(initialImageSelected(image));
},
[dispatch]
);
const recallAllParameters = useCallback(
(image: ImageDTO | undefined) => {
const type = image?.metadata?.type;
// not sure what this list should be
if (['t2l', 'l2l', 'inpaint'].includes(String(type))) {
dispatch(allParametersSet(image));
if (image?.metadata?.type === 'l2l') {
dispatch(setActiveTab('img2img'));
} else if (image?.metadata?.type === 't2l') {
dispatch(setActiveTab('txt2img'));
}
toaster({
title: t('toast.parametersSet'),
status: 'success',
duration: 2500,
isClosable: true,
});
} else {
toaster({
title: t('toast.parametersNotSet'),
description: t('toast.parametersNotSetDesc'),
status: 'error',
duration: 2500,
isClosable: true,
});
}
},
[t, toaster, dispatch]
);
return {
recallPrompt,
recallSeed,
recallInitialImage,
sendToImageToImage,
recallAllParameters,
};
};

View File

@ -1,23 +0,0 @@
import { getPromptAndNegative } from 'common/util/getPromptAndNegative';
import * as InvokeAI from 'app/types/invokeai';
import promptToString from 'common/util/promptToString';
import { useAppDispatch } from 'app/store/storeHooks';
import { setNegativePrompt, setPositivePrompt } from '../store/generationSlice';
import { useCallback } from 'react';
// TECHDEBT: We have two metadata prompt formats and need to handle recalling either of them.
// This hook provides a function to do that.
const useSetBothPrompts = () => {
const dispatch = useAppDispatch();
return useCallback(
(inputPrompt: InvokeAI.Prompt, negativePrompt: InvokeAI.Prompt) => {
dispatch(setPositivePrompt(inputPrompt));
dispatch(setNegativePrompt(negativePrompt));
},
[dispatch]
);
};
export default useSetBothPrompts;

View File

@ -0,0 +1,348 @@
import { useAppDispatch } from 'app/store/storeHooks';
import { useCallback } from 'react';
import { useTranslation } from 'react-i18next';
import {
modelSelected,
setCfgScale,
setHeight,
setImg2imgStrength,
setNegativePrompt,
setPositivePrompt,
setScheduler,
setSeed,
setSteps,
setWidth,
} from '../store/generationSlice';
import { isImageField } from 'services/types/guards';
import { initialImageSelected } from '../store/actions';
import { useAppToaster } from 'app/components/Toaster';
import { ImageDTO } from 'services/api';
import {
isValidCfgScale,
isValidHeight,
isValidModel,
isValidNegativePrompt,
isValidPositivePrompt,
isValidScheduler,
isValidSeed,
isValidSteps,
isValidStrength,
isValidWidth,
} from '../store/parameterZodSchemas';
export const useRecallParameters = () => {
const dispatch = useAppDispatch();
const toaster = useAppToaster();
const { t } = useTranslation();
const parameterSetToast = useCallback(() => {
toaster({
title: t('toast.parameterSet'),
status: 'info',
duration: 2500,
isClosable: true,
});
}, [t, toaster]);
const parameterNotSetToast = useCallback(() => {
toaster({
title: t('toast.parameterNotSet'),
status: 'warning',
duration: 2500,
isClosable: true,
});
}, [t, toaster]);
const allParameterSetToast = useCallback(() => {
toaster({
title: t('toast.parametersSet'),
status: 'info',
duration: 2500,
isClosable: true,
});
}, [t, toaster]);
const allParameterNotSetToast = useCallback(() => {
toaster({
title: t('toast.parametersNotSet'),
status: 'warning',
duration: 2500,
isClosable: true,
});
}, [t, toaster]);
/**
* Recall both prompts with toast
*/
const recallBothPrompts = useCallback(
(positivePrompt: unknown, negativePrompt: unknown) => {
if (
isValidPositivePrompt(positivePrompt) ||
isValidNegativePrompt(negativePrompt)
) {
if (isValidPositivePrompt(positivePrompt)) {
dispatch(setPositivePrompt(positivePrompt));
}
if (isValidNegativePrompt(negativePrompt)) {
dispatch(setNegativePrompt(negativePrompt));
}
parameterSetToast();
return;
}
parameterNotSetToast();
},
[dispatch, parameterSetToast, parameterNotSetToast]
);
/**
* Recall positive prompt with toast
*/
const recallPositivePrompt = useCallback(
(positivePrompt: unknown) => {
if (!isValidPositivePrompt(positivePrompt)) {
parameterNotSetToast();
return;
}
dispatch(setPositivePrompt(positivePrompt));
parameterSetToast();
},
[dispatch, parameterSetToast, parameterNotSetToast]
);
/**
* Recall negative prompt with toast
*/
const recallNegativePrompt = useCallback(
(negativePrompt: unknown) => {
if (!isValidNegativePrompt(negativePrompt)) {
parameterNotSetToast();
return;
}
dispatch(setNegativePrompt(negativePrompt));
parameterSetToast();
},
[dispatch, parameterSetToast, parameterNotSetToast]
);
/**
* Recall seed with toast
*/
const recallSeed = useCallback(
(seed: unknown) => {
if (!isValidSeed(seed)) {
parameterNotSetToast();
return;
}
dispatch(setSeed(seed));
parameterSetToast();
},
[dispatch, parameterSetToast, parameterNotSetToast]
);
/**
* Recall CFG scale with toast
*/
const recallCfgScale = useCallback(
(cfgScale: unknown) => {
if (!isValidCfgScale(cfgScale)) {
parameterNotSetToast();
return;
}
dispatch(setCfgScale(cfgScale));
parameterSetToast();
},
[dispatch, parameterSetToast, parameterNotSetToast]
);
/**
* Recall model with toast
*/
const recallModel = useCallback(
(model: unknown) => {
if (!isValidModel(model)) {
parameterNotSetToast();
return;
}
dispatch(modelSelected(model));
parameterSetToast();
},
[dispatch, parameterSetToast, parameterNotSetToast]
);
/**
* Recall scheduler with toast
*/
const recallScheduler = useCallback(
(scheduler: unknown) => {
if (!isValidScheduler(scheduler)) {
parameterNotSetToast();
return;
}
dispatch(setScheduler(scheduler));
parameterSetToast();
},
[dispatch, parameterSetToast, parameterNotSetToast]
);
/**
* Recall steps with toast
*/
const recallSteps = useCallback(
(steps: unknown) => {
if (!isValidSteps(steps)) {
parameterNotSetToast();
return;
}
dispatch(setSteps(steps));
parameterSetToast();
},
[dispatch, parameterSetToast, parameterNotSetToast]
);
/**
* Recall width with toast
*/
const recallWidth = useCallback(
(width: unknown) => {
if (!isValidWidth(width)) {
parameterNotSetToast();
return;
}
dispatch(setWidth(width));
parameterSetToast();
},
[dispatch, parameterSetToast, parameterNotSetToast]
);
/**
* Recall height with toast
*/
const recallHeight = useCallback(
(height: unknown) => {
if (!isValidHeight(height)) {
parameterNotSetToast();
return;
}
dispatch(setHeight(height));
parameterSetToast();
},
[dispatch, parameterSetToast, parameterNotSetToast]
);
/**
* Recall strength with toast
*/
const recallStrength = useCallback(
(strength: unknown) => {
if (!isValidStrength(strength)) {
parameterNotSetToast();
return;
}
dispatch(setImg2imgStrength(strength));
parameterSetToast();
},
[dispatch, parameterSetToast, parameterNotSetToast]
);
/**
* Sets initial image with toast
*/
const recallInitialImage = useCallback(
async (image: unknown) => {
if (!isImageField(image)) {
parameterNotSetToast();
return;
}
dispatch(initialImageSelected(image.image_name));
parameterSetToast();
},
[dispatch, parameterSetToast, parameterNotSetToast]
);
/**
* Sets image as initial image with toast
*/
const sendToImageToImage = useCallback(
(image: ImageDTO) => {
dispatch(initialImageSelected(image));
},
[dispatch]
);
const recallAllParameters = useCallback(
(image: ImageDTO | undefined) => {
if (!image || !image.metadata) {
allParameterNotSetToast();
return;
}
const {
cfg_scale,
height,
model,
positive_conditioning,
negative_conditioning,
scheduler,
seed,
steps,
width,
strength,
clip,
extra,
latents,
unet,
vae,
} = image.metadata;
if (isValidCfgScale(cfg_scale)) {
dispatch(setCfgScale(cfg_scale));
}
if (isValidModel(model)) {
dispatch(modelSelected(model));
}
if (isValidPositivePrompt(positive_conditioning)) {
dispatch(setPositivePrompt(positive_conditioning));
}
if (isValidNegativePrompt(negative_conditioning)) {
dispatch(setNegativePrompt(negative_conditioning));
}
if (isValidScheduler(scheduler)) {
dispatch(setScheduler(scheduler));
}
if (isValidSeed(seed)) {
dispatch(setSeed(seed));
}
if (isValidSteps(steps)) {
dispatch(setSteps(steps));
}
if (isValidWidth(width)) {
dispatch(setWidth(width));
}
if (isValidHeight(height)) {
dispatch(setHeight(height));
}
if (isValidStrength(strength)) {
dispatch(setImg2imgStrength(strength));
}
allParameterSetToast();
},
[allParameterNotSetToast, allParameterSetToast, dispatch]
);
return {
recallBothPrompts,
recallPositivePrompt,
recallNegativePrompt,
recallSeed,
recallInitialImage,
recallCfgScale,
recallModel,
recallScheduler,
recallSteps,
recallWidth,
recallHeight,
recallStrength,
recallAllParameters,
sendToImageToImage,
};
};

View File

@ -1,44 +1,53 @@
import type { PayloadAction } from '@reduxjs/toolkit';
import { createSlice } from '@reduxjs/toolkit';
import * as InvokeAI from 'app/types/invokeai';
import promptToString from 'common/util/promptToString';
import { clamp, sortBy } from 'lodash-es';
import { setAllParametersReducer } from './setAllParametersReducer';
import { receivedModels } from 'services/thunks/model';
import { Scheduler } from 'app/constants';
import { ImageDTO } from 'services/api';
import { configChanged } from 'features/system/store/configSlice';
import {
CfgScaleParam,
HeightParam,
ModelParam,
NegativePromptParam,
PositivePromptParam,
SchedulerParam,
SeedParam,
StepsParam,
StrengthParam,
WidthParam,
} from './parameterZodSchemas';
export interface GenerationState {
cfgScale: number;
height: number;
img2imgStrength: number;
cfgScale: CfgScaleParam;
height: HeightParam;
img2imgStrength: StrengthParam;
infillMethod: string;
initialImage?: ImageDTO;
iterations: number;
perlin: number;
positivePrompt: string;
negativePrompt: string;
scheduler: Scheduler;
positivePrompt: PositivePromptParam;
negativePrompt: NegativePromptParam;
scheduler: SchedulerParam;
seamBlur: number;
seamSize: number;
seamSteps: number;
seamStrength: number;
seed: number;
seed: SeedParam;
seedWeights: string;
shouldFitToWidthHeight: boolean;
shouldGenerateVariations: boolean;
shouldRandomizeSeed: boolean;
shouldUseNoiseSettings: boolean;
steps: number;
steps: StepsParam;
threshold: number;
tileSize: number;
variationAmount: number;
width: number;
width: WidthParam;
shouldUseSymmetry: boolean;
horizontalSymmetrySteps: number;
verticalSymmetrySteps: number;
model: string;
model: ModelParam;
shouldUseSeamless: boolean;
seamlessXAxis: boolean;
seamlessYAxis: boolean;
@ -84,27 +93,11 @@ export const generationSlice = createSlice({
name: 'generation',
initialState,
reducers: {
setPositivePrompt: (
state,
action: PayloadAction<string | InvokeAI.Prompt>
) => {
const newPrompt = action.payload;
if (typeof newPrompt === 'string') {
state.positivePrompt = newPrompt;
} else {
state.positivePrompt = promptToString(newPrompt);
}
setPositivePrompt: (state, action: PayloadAction<string>) => {
state.positivePrompt = action.payload;
},
setNegativePrompt: (
state,
action: PayloadAction<string | InvokeAI.Prompt>
) => {
const newPrompt = action.payload;
if (typeof newPrompt === 'string') {
state.negativePrompt = newPrompt;
} else {
state.negativePrompt = promptToString(newPrompt);
}
setNegativePrompt: (state, action: PayloadAction<string>) => {
state.negativePrompt = action.payload;
},
setIterations: (state, action: PayloadAction<number>) => {
state.iterations = action.payload;
@ -175,7 +168,6 @@ export const generationSlice = createSlice({
state.shouldGenerateVariations = true;
state.variationAmount = 0;
},
allParametersSet: setAllParametersReducer,
resetParametersState: (state) => {
return {
...state,
@ -279,7 +271,6 @@ export const {
setSeamless,
setSeamlessXAxis,
setSeamlessYAxis,
allParametersSet,
} = generationSlice.actions;
export default generationSlice.reducer;

View File

@ -0,0 +1,156 @@
import { NUMPY_RAND_MAX, SCHEDULERS } from 'app/constants';
import { z } from 'zod';
/**
* These zod schemas should match the pydantic node schemas.
*
* Parameters only need schemas if we want to recall them from metadata.
*
* Each parameter needs:
* - a zod schema
* - a type alias, inferred from the zod schema
* - a combo validation/type guard function, which returns true if the value is valid
*/
/**
* Zod schema for positive prompt parameter
*/
export const zPositivePrompt = z.string();
/**
* Type alias for positive prompt parameter, inferred from its zod schema
*/
export type PositivePromptParam = z.infer<typeof zPositivePrompt>;
/**
* Validates/type-guards a value as a positive prompt parameter
*/
export const isValidPositivePrompt = (
val: unknown
): val is PositivePromptParam => zPositivePrompt.safeParse(val).success;
/**
* Zod schema for negative prompt parameter
*/
export const zNegativePrompt = z.string();
/**
* Type alias for negative prompt parameter, inferred from its zod schema
*/
export type NegativePromptParam = z.infer<typeof zNegativePrompt>;
/**
* Validates/type-guards a value as a negative prompt parameter
*/
export const isValidNegativePrompt = (
val: unknown
): val is NegativePromptParam => zNegativePrompt.safeParse(val).success;
/**
* Zod schema for steps parameter
*/
export const zSteps = z.number().int().min(1);
/**
* Type alias for steps parameter, inferred from its zod schema
*/
export type StepsParam = z.infer<typeof zSteps>;
/**
* Validates/type-guards a value as a steps parameter
*/
export const isValidSteps = (val: unknown): val is StepsParam =>
zSteps.safeParse(val).success;
/**
* Zod schema for CFG scale parameter
*/
export const zCfgScale = z.number().min(1);
/**
* Type alias for CFG scale parameter, inferred from its zod schema
*/
export type CfgScaleParam = z.infer<typeof zCfgScale>;
/**
* Validates/type-guards a value as a CFG scale parameter
*/
export const isValidCfgScale = (val: unknown): val is CfgScaleParam =>
zCfgScale.safeParse(val).success;
/**
* Zod schema for scheduler parameter
*/
export const zScheduler = z.enum(SCHEDULERS);
/**
* Type alias for scheduler parameter, inferred from its zod schema
*/
export type SchedulerParam = z.infer<typeof zScheduler>;
/**
* Validates/type-guards a value as a scheduler parameter
*/
export const isValidScheduler = (val: unknown): val is SchedulerParam =>
zScheduler.safeParse(val).success;
/**
* Zod schema for seed parameter
*/
export const zSeed = z.number().int().min(0).max(NUMPY_RAND_MAX);
/**
* Type alias for seed parameter, inferred from its zod schema
*/
export type SeedParam = z.infer<typeof zSeed>;
/**
* Validates/type-guards a value as a seed parameter
*/
export const isValidSeed = (val: unknown): val is SeedParam =>
zSeed.safeParse(val).success;
/**
* Zod schema for width parameter
*/
export const zWidth = z.number().multipleOf(8).min(64);
/**
* Type alias for width parameter, inferred from its zod schema
*/
export type WidthParam = z.infer<typeof zWidth>;
/**
* Validates/type-guards a value as a width parameter
*/
export const isValidWidth = (val: unknown): val is WidthParam =>
zWidth.safeParse(val).success;
/**
* Zod schema for height parameter
*/
export const zHeight = z.number().multipleOf(8).min(64);
/**
* Type alias for height parameter, inferred from its zod schema
*/
export type HeightParam = z.infer<typeof zHeight>;
/**
* Validates/type-guards a value as a height parameter
*/
export const isValidHeight = (val: unknown): val is HeightParam =>
zHeight.safeParse(val).success;
/**
* Zod schema for model parameter
* TODO: Make this a dynamically generated enum?
*/
export const zModel = z.string();
/**
* Type alias for model parameter, inferred from its zod schema
*/
export type ModelParam = z.infer<typeof zModel>;
/**
* Validates/type-guards a value as a model parameter
*/
export const isValidModel = (val: unknown): val is ModelParam =>
zModel.safeParse(val).success;
/**
* Zod schema for l2l strength parameter
*/
export const zStrength = z.number().min(0).max(1);
/**
* Type alias for l2l strength parameter, inferred from its zod schema
*/
export type StrengthParam = z.infer<typeof zStrength>;
/**
* Validates/type-guards a value as a l2l strength parameter
*/
export const isValidStrength = (val: unknown): val is StrengthParam =>
zStrength.safeParse(val).success;

View File

@ -1,77 +0,0 @@
import { Draft, PayloadAction } from '@reduxjs/toolkit';
import { GenerationState } from './generationSlice';
import { ImageDTO, ImageToImageInvocation } from 'services/api';
import { isScheduler } from 'app/constants';
export const setAllParametersReducer = (
state: Draft<GenerationState>,
action: PayloadAction<ImageDTO | undefined>
) => {
const metadata = action.payload?.metadata;
if (!metadata) {
return;
}
// not sure what this list should be
if (
metadata.type === 't2l' ||
metadata.type === 'l2l' ||
metadata.type === 'inpaint'
) {
const {
cfg_scale,
height,
model,
positive_conditioning,
negative_conditioning,
scheduler,
seed,
steps,
width,
} = metadata;
if (cfg_scale !== undefined) {
state.cfgScale = Number(cfg_scale);
}
if (height !== undefined) {
state.height = Number(height);
}
if (model !== undefined) {
state.model = String(model);
}
if (positive_conditioning !== undefined) {
state.positivePrompt = String(positive_conditioning);
}
if (negative_conditioning !== undefined) {
state.negativePrompt = String(negative_conditioning);
}
if (scheduler !== undefined) {
const schedulerString = String(scheduler);
if (isScheduler(schedulerString)) {
state.scheduler = schedulerString;
}
}
if (seed !== undefined) {
state.seed = Number(seed);
state.shouldRandomizeSeed = false;
}
if (steps !== undefined) {
state.steps = Number(steps);
}
if (width !== undefined) {
state.width = Number(width);
}
}
if (metadata.type === 'l2l') {
const { fit, image } = metadata as ImageToImageInvocation;
if (fit !== undefined) {
state.shouldFitToWidthHeight = Boolean(fit);
}
// if (image !== undefined) {
// state.initialImage = image;
// }
}
};