Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(ui): improve scheduler selection logic

Browse Source 
- remove UI-specific state (the enabled schedulers) from redux, instead derive it in a selector
- simplify logic by putting schedulers in an object instead of an array
- rename `activeSchedulers` to `enabledSchedulers`
- remove need for `useEffect()` when `enabledSchedulers` changes by adding a listener for the `enabledSchedulersChanged` action/event to `generationSlice`
- increase type safety by making `enabledSchedulers` an array of `SchedulerParam`, which is created by the zod schema for scheduler
This commit is contained in:
psychedelicious 2023-06-18 19:34:37 +10:00
parent 150059f704
commit 94cfcdc411
6 changed files with 127 additions and 97 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

63
invokeai/frontend/web/src/app/constants.ts
View File

@ -1,27 +1,8 @@
import { SelectItem } from '@mantine/core';
export const SCHEDULERS: SelectItem[] = [
{ label: 'Euler', value: 'euler', group: 'Standard' },
{ label: 'DEIS', value: 'deis', group: 'Standard' },
{ label: 'DDIM', value: 'ddim', group: 'Standard' },
{ label: 'DDPM', value: 'ddpm', group: 'Standard' },
{ label: 'DPM++ 2S', value: 'dpmpp_2s', group: 'Standard' },
{ label: 'DPM++ 2M', value: 'dpmpp_2m', group: 'Standard' },
{ label: 'Heun', value: 'heun', group: 'Standard' },
{ label: 'KDPM 2', value: 'kdpm_2', group: 'Standard' },
{ label: 'LMS', value: 'lms', group: 'Standard' },
{ label: 'PNDM', value: 'pndm', group: 'Standard' },
{ label: 'UniPC', value: 'unipc', group: 'Standard' },
{ label: 'Euler Karras', value: 'euler_k', group: 'Karras' },
{ label: 'DPM++ 2S Karras', value: 'dpmpp_2s_k', group: 'Karras' },
{ label: 'DPM++ 2M Karras', value: 'dpmpp_2m_k', group: 'Karras' },
{ label: 'Heun Karras', value: 'heun_k', group: 'Karras' },
{ label: 'LMS Karras', value: 'lms_k', group: 'Karras' },
{ label: 'Euler Ancestral', value: 'euler_a', group: 'Ancestral' },
{ label: 'KDPM 2 Ancestral', value: 'kdpm_2_a', group: 'Ancestral' },
];
import { SchedulerParam } from 'features/parameters/store/parameterZodSchemas';
// zod needs the array to be `as const` to infer the type correctly
// this is the source of the `SchedulerParam` type, which is generated by zod
export const SCHEDULER_NAMES_AS_CONST = [
'euler',
'deis',
@ -43,7 +24,45 @@ export const SCHEDULER_NAMES_AS_CONST = [
'kdpm_2_a',
] as const;
export const SCHEDULER_NAMES = [...SCHEDULER_NAMES_AS_CONST];
export const DEFAULT_SCHEDULER_NAME = 'euler';
export const SCHEDULER_NAMES: SchedulerParam[] = [...SCHEDULER_NAMES_AS_CONST];
export const SCHEDULER_SELECT_ITEMS: Record<
(typeof SCHEDULER_NAMES)[number],
SelectItem & { label: string; value: SchedulerParam; group: string }
> = {
euler: { label: 'Euler', value: 'euler', group: 'Standard' },
deis: { label: 'DEIS', value: 'deis', group: 'Standard' },
ddim: { label: 'DDIM', value: 'ddim', group: 'Standard' },
ddpm: { label: 'DDPM', value: 'ddpm', group: 'Standard' },
dpmpp_2s: { label: 'DPM++ 2S', value: 'dpmpp_2s', group: 'Standard' },
dpmpp_2m: { label: 'DPM++ 2M', value: 'dpmpp_2m', group: 'Standard' },
heun: { label: 'Heun', value: 'heun', group: 'Standard' },
kdpm_2: { label: 'KDPM 2', value: 'kdpm_2', group: 'Standard' },
lms: { label: 'LMS', value: 'lms', group: 'Standard' },
pndm: { label: 'PNDM', value: 'pndm', group: 'Standard' },
unipc: { label: 'UniPC', value: 'unipc', group: 'Standard' },
euler_k: { label: 'Euler Karras', value: 'euler_k', group: 'Karras' },
dpmpp_2s_k: {
label: 'DPM++ 2S Karras',
value: 'dpmpp_2s_k',
group: 'Karras',
},
dpmpp_2m_k: {
label: 'DPM++ 2M Karras',
value: 'dpmpp_2m_k',
group: 'Karras',
},
heun_k: { label: 'Heun Karras', value: 'heun_k', group: 'Karras' },
lms_k: { label: 'LMS Karras', value: 'lms_k', group: 'Karras' },
euler_a: { label: 'Euler Ancestral', value: 'euler_a', group: 'Ancestral' },
kdpm_2_a: {
label: 'KDPM 2 Ancestral',
value: 'kdpm_2_a',
group: 'Ancestral',
},
};
export type Scheduler = (typeof SCHEDULER_NAMES)[number];

63
invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamScheduler.tsx
View File

@ -1,48 +1,47 @@
import { SCHEDULER_NAMES, Scheduler } from 'app/constants';
import { RootState } from 'app/store/store';
import { createSelector } from '@reduxjs/toolkit';
import { SCHEDULER_SELECT_ITEMS } from 'app/constants';
import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
import IAIMantineSelect from 'common/components/IAIMantineSelect';
import { generationSelector } from 'features/parameters/store/generationSelectors';
import { setScheduler } from 'features/parameters/store/generationSlice';
import { setSelectedSchedulers } from 'features/ui/store/uiSlice';
import { memo, useCallback, useEffect } from 'react';
import { SchedulerParam } from 'features/parameters/store/parameterZodSchemas';
import { uiSelector } from 'features/ui/store/uiSelectors';
import { memo, useCallback } from 'react';
import { useTranslation } from 'react-i18next';
const selector = createSelector(
[uiSelector, generationSelector],
(ui, generation) => {
const { scheduler } = generation;
const { enabledSchedulers } = ui;
const data = enabledSchedulers
.map(
(schedulerName) =>
SCHEDULER_SELECT_ITEMS[schedulerName as SchedulerParam]
)
.sort((a, b) => a.label.localeCompare(b.label));
return {
scheduler,
data,
};
},
defaultSelectorOptions
);
const ParamScheduler = () => {
const scheduler = useAppSelector(
(state: RootState) => state.generation.scheduler
);
const selectedSchedulers = useAppSelector(
(state: RootState) => state.ui.selectedSchedulers
);
const activeSchedulers = useAppSelector(
(state: RootState) => state.ui.activeSchedulers
);
const dispatch = useAppDispatch();
const { t } = useTranslation();
useEffect(() => {
if (selectedSchedulers.length === 0) {
dispatch(setSelectedSchedulers(SCHEDULER_NAMES));
}
const schedulerFound = activeSchedulers.find(
(activeSchedulers) => activeSchedulers.value === scheduler
);
if (!schedulerFound) {
dispatch(setScheduler(activeSchedulers[0].value as Scheduler));
}
}, [dispatch, selectedSchedulers, scheduler, activeSchedulers]);
const { scheduler, data } = useAppSelector(selector);
const handleChange = useCallback(
(v: string | null) => {
if (!v) {
return;
}
dispatch(setScheduler(v as Scheduler));
dispatch(setScheduler(v as SchedulerParam));
},
[dispatch]
);
@ -51,7 +50,7 @@ const ParamScheduler = () => {
<IAIMantineSelect
label={t('parameters.scheduler')}
value={scheduler}
data={activeSchedulers}
data={data}
onChange={handleChange}
/>
);

22
invokeai/frontend/web/src/features/parameters/store/generationSlice.ts
View File

@ -1,6 +1,5 @@
import type { PayloadAction } from '@reduxjs/toolkit';
import { createSlice } from '@reduxjs/toolkit';
import { Scheduler } from 'app/constants';
import { configChanged } from 'features/system/store/configSlice';
import { clamp, sortBy } from 'lodash-es';
import { ImageDTO } from 'services/api';
@ -18,6 +17,8 @@ import {
StrengthParam,
WidthParam,
} from './parameterZodSchemas';
import { enabledSchedulersChanged } from 'features/ui/store/uiSlice';
import { DEFAULT_SCHEDULER_NAME } from 'app/constants';
export interface GenerationState {
cfgScale: CfgScaleParam;
@ -63,7 +64,7 @@ export const initialGenerationState: GenerationState = {
perlin: 0,
positivePrompt: '',
negativePrompt: '',
scheduler: 'euler',
scheduler: DEFAULT_SCHEDULER_NAME,
seamBlur: 16,
seamSize: 96,
seamSteps: 30,
@ -133,7 +134,7 @@ export const generationSlice = createSlice({
setWidth: (state, action: PayloadAction<number>) => {
state.width = action.payload;
},
setScheduler: (state, action: PayloadAction<Scheduler>) => {
setScheduler: (state, action: PayloadAction<SchedulerParam>) => {
state.scheduler = action.payload;
},
setSeed: (state, action: PayloadAction<number>) => {
@ -241,6 +242,21 @@ export const generationSlice = createSlice({
state.initialImage.thumbnail_url = thumbnail_url;
}
});
builder.addCase(enabledSchedulersChanged, (state, action) => {
const enabledSchedulers = action.payload;
if (!action.payload.length) {
// This means the user cleared the enabled schedulers multi-select. We need to set the scheduler to the default
state.scheduler = DEFAULT_SCHEDULER_NAME;
return;
}
if (!enabledSchedulers.includes(state.scheduler)) {
// The current scheduler is now disabled, change it to the first enabled one
state.scheduler = action.payload[0];
}
});
},
});

35
invokeai/frontend/web/src/features/system/components/SettingsModal/SettingsSchedulers.tsx
View File

@ -1,30 +1,39 @@
import { SCHEDULERS } from 'app/constants';
import { SCHEDULER_SELECT_ITEMS } from 'app/constants';
import { RootState } from 'app/store/store';
import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
import IAIMantineMultiSelect from 'common/components/IAIMantineMultiSelect';
import { setSelectedSchedulers } from 'features/ui/store/uiSlice';
import { SchedulerParam } from 'features/parameters/store/parameterZodSchemas';
import { enabledSchedulersChanged } from 'features/ui/store/uiSlice';
import { map } from 'lodash-es';
import { useCallback } from 'react';
import { useTranslation } from 'react-i18next';
const data = map(SCHEDULER_SELECT_ITEMS).sort((a, b) =>
a.label.localeCompare(b.label)
);
export default function SettingsSchedulers() {
const dispatch = useAppDispatch();
const selectedSchedulers = useAppSelector(
(state: RootState) => state.ui.selectedSchedulers
);
const { t } = useTranslation();
const schedulerSettingsHandler = (v: string[]) => {
dispatch(setSelectedSchedulers(v));
};
const enabledSchedulers = useAppSelector(
(state: RootState) => state.ui.enabledSchedulers
);
const handleChange = useCallback(
(v: string[]) => {
dispatch(enabledSchedulersChanged(v as SchedulerParam[]));
},
[dispatch]
);
return (
<IAIMantineMultiSelect
label={t('settings.availableSchedulers')}
value={selectedSchedulers}
data={SCHEDULERS}
onChange={schedulerSettingsHandler}
value={enabledSchedulers}
data={data}
onChange={handleChange}
clearable
searchable
maxSelectedValues={99}

36
invokeai/frontend/web/src/features/ui/store/uiSlice.ts
View File

@ -1,11 +1,11 @@
import { SelectItem } from '@mantine/core';
import type { PayloadAction } from '@reduxjs/toolkit';
import { createSlice } from '@reduxjs/toolkit';
import { SCHEDULERS } from 'app/constants';
import { initialImageChanged } from 'features/parameters/store/generationSlice';
import { setActiveTabReducer } from './extraReducers';
import { InvokeTabName } from './tabMap';
import { AddNewModelType, UIState } from './uiTypes';
import { SchedulerParam } from 'features/parameters/store/parameterZodSchemas';
import { DEFAULT_SCHEDULER_NAME, SCHEDULER_NAMES } from 'app/constants';
export const initialUIState: UIState = {
activeTab: 0,
@ -21,8 +21,7 @@ export const initialUIState: UIState = {
shouldShowGallery: true,
shouldHidePreview: false,
shouldShowProgressInViewer: true,
activeSchedulers: [],
selectedSchedulers: [],
enabledSchedulers: [],
};
export const uiSlice = createSlice({
@ -96,27 +95,16 @@ export const uiSlice = createSlice({
setShouldShowProgressInViewer: (state, action: PayloadAction<boolean>) => {
state.shouldShowProgressInViewer = action.payload;
},
setSelectedSchedulers: (state, action: PayloadAction<string[]>) => {
const selectedSchedulerData: SelectItem[] = [];
let selectedSchedulers = [...action.payload];
if (selectedSchedulers.length === 0) {
selectedSchedulers = [SCHEDULERS[0].value];
enabledSchedulersChanged: (
state,
action: PayloadAction<SchedulerParam[]>
) => {
if (action.payload.length === 0) {
state.enabledSchedulers = [DEFAULT_SCHEDULER_NAME];
return;
}
selectedSchedulers.forEach((item) => {
const schedulerData = SCHEDULERS.find(
(scheduler) => scheduler.value === item
);
if (schedulerData) {
selectedSchedulerData.push(schedulerData);
}
});
state.activeSchedulers = selectedSchedulerData;
state.selectedSchedulers = action.payload;
state.enabledSchedulers = action.payload;
},
},
extraReducers(builder) {
@ -144,7 +132,7 @@ export const {
toggleParametersPanel,
toggleGalleryPanel,
setShouldShowProgressInViewer,
setSelectedSchedulers,
enabledSchedulersChanged,
} = uiSlice.actions;
export default uiSlice.reducer;

5
invokeai/frontend/web/src/features/ui/store/uiTypes.ts
View File

@ -1,4 +1,4 @@
import { SelectItem } from '@mantine/core';
import { SchedulerParam } from 'features/parameters/store/parameterZodSchemas';
export type AddNewModelType = 'ckpt' | 'diffusers' | null;
@ -28,6 +28,5 @@ export interface UIState {
shouldPinGallery: boolean;
shouldShowGallery: boolean;
shouldShowProgressInViewer: boolean;
activeSchedulers: SelectItem[];
selectedSchedulers: string[];
enabledSchedulers: SchedulerParam[];
}