apextoaster
/
onnx-web
1
0
Fork 0
Code Actions

optimize zustand selectors

This commit is contained in:
Sean Sube 2023-07-21 22:45:20 -05:00
parent 626ca18d7f
commit ffc941548d
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
8 changed files with 135 additions and 77 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
gui/src/components/Profiles.tsx
View File

@ -21,15 +21,15 @@ import { useTranslation } from 'react-i18next';
import { useStore } from 'zustand';
import { BaseImgParams, HighresParams, ImageMetadata, Txt2ImgParams, UpscaleParams } from '../client/types.js';
import { StateContext } from '../state.js';
import { OnnxState, StateContext } from '../state.js';
import { DeepPartial } from '../types.js';
const { useState } = React;
export interface ProfilesProps {
highres: HighresParams;
params: BaseImgParams;
upscale: UpscaleParams;
selectHighres(state: OnnxState): HighresParams;
selectParams(state: OnnxState): BaseImgParams;
selectUpscale(state: OnnxState): UpscaleParams;
setHighres(params: Partial<HighresParams>): void;
setParams(params: Partial<BaseImgParams>): void;
@ -116,11 +116,12 @@ export function Profiles(props: ProfilesProps) {
<Button
variant='contained'
onClick={() => {
const innerState = state.getState();
saveProfile({
params: props.params,
params: props.selectParams(innerState),
name: profileName,
highres: props.highres,
upscale: props.upscale,
highres: props.selectHighres(innerState),
upscale: props.selectUpscale(innerState),
});
setDialogOpen(false);
setProfileName('');
@ -160,7 +161,12 @@ export function Profiles(props: ProfilesProps) {
/>
</Button>
<Button component='label' variant='contained' onClick={() => {
downloadParamsAsFile(props);
const innerState = state.getState();
downloadParamsAsFile({
params: props.selectParams(innerState),
highres: props.selectHighres(innerState),
upscale: props.selectUpscale(innerState),
});
}}>
<Download />
</Button>

10
gui/src/components/control/HighresControl.tsx
View File

@ -3,19 +3,23 @@ import { Checkbox, FormControl, FormControlLabel, InputLabel, MenuItem, Select,
import * as React from 'react';
import { useContext } from 'react';
import { useTranslation } from 'react-i18next';
import { useStore } from 'zustand';
import { HighresParams } from '../../client/types.js';
import { ConfigContext } from '../../state.js';
import { ConfigContext, OnnxState, StateContext } from '../../state.js';
import { NumericField } from '../input/NumericField.js';
export interface HighresControlProps {
highres: HighresParams;
selectHighres(state: OnnxState): HighresParams;
setHighres(params: Partial<HighresParams>): void;
}
export function HighresControl(props: HighresControlProps) {
// eslint-disable-next-line @typescript-eslint/unbound-method
const { highres, setHighres } = props;
const { selectHighres, setHighres } = props;
const state = mustExist(useContext(StateContext));
const highres = useStore(state, selectHighres);
const { params } = mustExist(useContext(ConfigContext));
const { t } = useTranslation();

10
gui/src/components/control/UpscaleControl.tsx
View File

@ -3,19 +3,23 @@ import { Checkbox, FormControl, FormControlLabel, InputLabel, MenuItem, Select,
import * as React from 'react';
import { useContext } from 'react';
import { useTranslation } from 'react-i18next';
import { useStore } from 'zustand';
import { UpscaleParams } from '../../client/types.js';
import { ConfigContext } from '../../state.js';
import { ConfigContext, OnnxState, StateContext } from '../../state.js';
import { NumericField } from '../input/NumericField.js';
export interface UpscaleControlProps {
upscale: UpscaleParams;
selectUpscale(state: OnnxState): UpscaleParams;
setUpscale(params: Partial<UpscaleParams>): void;
}
export function UpscaleControl(props: UpscaleControlProps) {
// eslint-disable-next-line @typescript-eslint/unbound-method
const { upscale, setUpscale } = props;
const { selectUpscale, setUpscale } = props;
const state = mustExist(useContext(StateContext));
const upscale = useStore(state, selectUpscale);
const { params } = mustExist(useContext(ConfigContext));
const { t } = useTranslation();

22
gui/src/components/tab/Blend.tsx
View File

@ -1,13 +1,14 @@
import { mustDefault, mustExist } from '@apextoaster/js-utils';
import { Box, Button, Stack } from '@mui/material';
import { useMutation, useQueryClient } from '@tanstack/react-query';
import * as React from 'react';
import { useContext } from 'react';
import { useTranslation } from 'react-i18next';
import { useMutation, useQueryClient } from '@tanstack/react-query';
import { useStore } from 'zustand';
import { BlendParams, ModelParams, UpscaleParams } from '../../client/types.js';
import { IMAGE_FILTER } from '../../config.js';
import { BLEND_SOURCES, ClientContext, StateContext } from '../../state.js';
import { BLEND_SOURCES, ClientContext, OnnxState, StateContext, TabState } from '../../state.js';
import { range } from '../../utils.js';
import { UpscaleControl } from '../control/UpscaleControl.js';
import { ImageInput } from '../input/ImageInput.js';
@ -33,8 +34,7 @@ export function Blend() {
const state = mustExist(useContext(StateContext));
const brush = useStore(state, (s) => s.blendBrush);
const blend = useStore(state, (s) => s.blend);
const upscale = useStore(state, (s) => s.blendUpscale);
const blend = useStore(state, selectParams);
// eslint-disable-next-line @typescript-eslint/unbound-method
const setBlend = useStore(state, (s) => s.setBlend);
// eslint-disable-next-line @typescript-eslint/unbound-method
@ -77,7 +77,7 @@ export function Blend() {
}}
setBrush={setBrush}
/>
<UpscaleControl upscale={upscale} setUpscale={setUpscale} />
<UpscaleControl selectUpscale={selectUpscale} setUpscale={setUpscale} />
<Button
disabled={sources.length < 2}
variant='contained'
@ -86,3 +86,15 @@ export function Blend() {
</Stack>
</Box>;
}
export function selectModel(state: OnnxState): ModelParams {
return state.blendModel;
}
export function selectParams(state: OnnxState): TabState<BlendParams> {
return state.blend;
}
export function selectUpscale(state: OnnxState): UpscaleParams {
return state.blendUpscale;
}

27
gui/src/components/tab/Img2Img.tsx
View File

@ -22,10 +22,13 @@ export function Img2Img() {
const { params } = mustExist(useContext(ConfigContext));
async function uploadSource() {
const innerState = state.getState();
const img2img = selectParams(innerState);
const { image, retry } = await client.img2img(model, {
...img2img,
source: mustExist(img2img.source), // TODO: show an error if this doesn't exist
}, upscale, highres);
}, selectUpscale(innerState), selectHighres(innerState));
pushHistory(image, retry);
}
@ -46,9 +49,9 @@ export function Img2Img() {
const state = mustExist(useContext(StateContext));
const model = useStore(state, selectModel);
const source = useStore(state, (s) => s.img2img.source);
const img2img = useStore(state, selectParams);
const highres = useStore(state, selectHighres);
const upscale = useStore(state, selectUpscale);
const sourceFilter = useStore(state, (s) => s.img2img.sourceFilter);
const strength = useStore(state, (s) => s.img2img.strength);
const loopback = useStore(state, (s) => s.img2img.loopback);
// eslint-disable-next-line @typescript-eslint/unbound-method
const setImg2Img = useStore(state, (s) => s.setImg2Img);
// eslint-disable-next-line @typescript-eslint/unbound-method
@ -64,11 +67,11 @@ export function Img2Img() {
return <Box>
<Stack spacing={2}>
<Profiles
params={img2img}
selectHighres={selectHighres}
selectParams={selectParams}
selectUpscale={selectUpscale}
setParams={setImg2Img}
highres={highres}
setHighres={setHighres}
upscale={upscale}
setUpscale={setUpscale}
/>
<ModelControl model={model} setModel={setModel} />
@ -108,7 +111,7 @@ export function Img2Img() {
selector: (f) => f.source,
}}
showNone
value={img2img.sourceFilter}
value={sourceFilter}
onChange={(newFilter) => {
setImg2Img({
sourceFilter: newFilter,
@ -121,7 +124,7 @@ export function Img2Img() {
min={params.strength.min}
max={params.strength.max}
step={params.strength.step}
value={img2img.strength}
value={strength}
onChange={(value) => {
setImg2Img({
strength: value,
@ -133,7 +136,7 @@ export function Img2Img() {
min={params.loopback.min}
max={params.loopback.max}
step={params.loopback.step}
value={img2img.loopback}
value={loopback}
onChange={(value) => {
setImg2Img({
loopback: value,
@ -141,8 +144,8 @@ export function Img2Img() {
}}
/>
</Stack>
<HighresControl highres={highres} setHighres={setHighres} />
<UpscaleControl upscale={upscale} setUpscale={setUpscale} />
<HighresControl selectHighres={selectHighres} setHighres={setHighres} />
<UpscaleControl selectUpscale={selectUpscale} setUpscale={setUpscale} />
<Button
disabled={doesExist(source) === false}
variant='contained'

57
gui/src/components/tab/Inpaint.tsx
View File

@ -6,6 +6,7 @@ import { useContext } from 'react';
import { useTranslation } from 'react-i18next';
import { useStore } from 'zustand';
import { HighresParams, InpaintParams, ModelParams, UpscaleParams } from '../../client/types.js';
import { IMAGE_FILTER, STALE_TIME } from '../../config.js';
import { ClientContext, ConfigContext, OnnxState, StateContext, TabState } from '../../state.js';
import { HighresControl } from '../control/HighresControl.js';
@ -18,7 +19,6 @@ import { MaskCanvas } from '../input/MaskCanvas.js';
import { NumericField } from '../input/NumericField.js';
import { QueryList } from '../input/QueryList.js';
import { Profiles } from '../Profiles.js';
import { ModelParams, InpaintParams, HighresParams, UpscaleParams } from '../../client/types.js';
export function Inpaint() {
const { params } = mustExist(useContext(ConfigContext));
@ -32,28 +32,31 @@ export function Inpaint() {
});
async function uploadSource(): Promise<void> {
const innerState = state.getState();
const inpaint = selectParams(innerState);
if (outpaint.enabled) {
const { image, retry } = await client.outpaint(model, {
...inpaint,
...outpaint,
mask: mustExist(inpaint.mask),
source: mustExist(inpaint.source),
}, upscale, highres);
mask: mustExist(mask),
source: mustExist(source),
}, selectUpscale(innerState), selectHighres(innerState));
pushHistory(image, retry);
} else {
const { image, retry } = await client.inpaint(model, {
...inpaint,
mask: mustExist(inpaint.mask),
source: mustExist(inpaint.source),
}, upscale, highres);
mask: mustExist(mask),
source: mustExist(source),
}, selectUpscale(innerState), selectHighres(innerState));
pushHistory(image, retry);
}
}
function preventInpaint(): boolean {
return doesExist(inpaint.source) === false || doesExist(inpaint.mask) === false;
return doesExist(source) === false || doesExist(mask) === false;
}
function supportsInpaint(): boolean {
@ -61,10 +64,14 @@ export function Inpaint() {
}
const state = mustExist(useContext(StateContext));
const inpaint = useStore(state, selectParams);
const highres = useStore(state, selectHighres);
const source = useStore(state, (s) => s.inpaint.source);
const mask = useStore(state, (s) => s.inpaint.mask);
const strength = useStore(state, (s) => s.inpaint.strength);
const noise = useStore(state, (s) => s.inpaint.noise);
const filter = useStore(state, (s) => s.inpaint.filter);
const tileOrder = useStore(state, (s) => s.inpaint.tileOrder);
const fillColor = useStore(state, (s) => s.inpaint.fillColor);
const model = useStore(state, selectModel);
const upscale = useStore(state, selectUpscale);
const outpaint = useStore(state, (s) => s.outpaint);
const brush = useStore(state, (s) => s.inpaintBrush);
@ -99,18 +106,18 @@ export function Inpaint() {
return <Box>
<Stack spacing={2}>
<Profiles
params={inpaint}
selectHighres={selectHighres}
selectParams={selectParams}
selectUpscale={selectUpscale}
setParams={setInpaint}
highres={highres}
setHighres={setHighres}
upscale={upscale}
setUpscale={setUpscale}
/>
<ModelControl model={model} setModel={setModel} />
{renderBanner()}
<ImageInput
filter={IMAGE_FILTER}
image={inpaint.source}
image={source}
label={t('input.image.source')}
hideSelection={true}
onChange={(file) => {
@ -121,7 +128,7 @@ export function Inpaint() {
/>
<ImageInput
filter={IMAGE_FILTER}
image={inpaint.mask}
image={mask}
label={t('input.image.mask')}
hideSelection={true}
onChange={(file) => {
@ -132,8 +139,8 @@ export function Inpaint() {
/>
<MaskCanvas
brush={brush}
source={inpaint.source}
mask={inpaint.mask}
source={source}
mask={mask}
onSave={(file) => {
setInpaint({
mask: file,
@ -152,7 +159,7 @@ export function Inpaint() {
min={params.strength.min}
max={params.strength.max}
step={params.strength.step}
value={inpaint.strength}
value={strength}
onChange={(value) => {
setInpaint({
strength: value,
@ -168,7 +175,7 @@ export function Inpaint() {
result: filters,
selector: (f) => f.mask,
}}
value={inpaint.filter}
value={filter}
onChange={(newFilter) => {
setInpaint({
filter: newFilter,
@ -182,7 +189,7 @@ export function Inpaint() {
query={{
result: noises,
}}
value={inpaint.noise}
value={noise}
onChange={(newNoise) => {
setInpaint({
noise: newNoise,
@ -194,7 +201,7 @@ export function Inpaint() {
<Select
labelId={'outpaint-tiling'}
label={t('parameter.tileOrder')}
value={inpaint.tileOrder}
value={tileOrder}
onChange={(e) => {
setInpaint({
tileOrder: e.target.value,
@ -212,7 +219,7 @@ export function Inpaint() {
sx={{ mx: 1 }}
control={
<input
defaultValue={inpaint.fillColor}
defaultValue={fillColor}
name='fill-color'
type='color'
onBlur={(event) => {
@ -226,8 +233,8 @@ export function Inpaint() {
</Stack>
</Stack>
<OutpaintControl />
<HighresControl highres={highres} setHighres={setHighres} />
<UpscaleControl upscale={upscale} setUpscale={setUpscale} />
<HighresControl selectHighres={selectHighres} setHighres={setHighres} />
<UpscaleControl selectUpscale={selectUpscale} setUpscale={setUpscale} />
<Button
disabled={preventInpaint()}
variant='contained'

30
gui/src/components/tab/Txt2Img.tsx
View File

@ -1,25 +1,26 @@
import { mustExist } from '@apextoaster/js-utils';
import { Box, Button, Stack } from '@mui/material';
import { useMutation, useQueryClient } from '@tanstack/react-query';
import * as React from 'react';
import { useContext } from 'react';
import { useTranslation } from 'react-i18next';
import { useMutation, useQueryClient } from '@tanstack/react-query';
import { useStore } from 'zustand';
import { HighresParams, ModelParams, Txt2ImgParams, UpscaleParams } from '../../client/types.js';
import { ClientContext, ConfigContext, OnnxState, StateContext, TabState } from '../../state.js';
import { HighresControl } from '../control/HighresControl.js';
import { ImageControl } from '../control/ImageControl.js';
import { ModelControl } from '../control/ModelControl.js';
import { UpscaleControl } from '../control/UpscaleControl.js';
import { NumericField } from '../input/NumericField.js';
import { ModelControl } from '../control/ModelControl.js';
import { Profiles } from '../Profiles.js';
import { HighresParams, ModelParams, Txt2ImgParams, UpscaleParams } from '../../client/types.js';
export function Txt2Img() {
const { params } = mustExist(useContext(ConfigContext));
async function generateImage() {
const { image, retry } = await client.txt2img(model, txt2img, upscale, highres);
const innerState = state.getState();
const { image, retry } = await client.txt2img(model, selectParams(innerState), selectUpscale(innerState), selectHighres(innerState));
pushHistory(image, retry);
}
@ -31,10 +32,9 @@ export function Txt2Img() {
});
const state = mustExist(useContext(StateContext));
const txt2img = useStore(state, selectParams);
const height = useStore(state, (s) => s.txt2img.height);
const width = useStore(state, (s) => s.txt2img.width);
const model = useStore(state, selectModel);
const highres = useStore(state, selectHighres);
const upscale = useStore(state, selectUpscale);
// eslint-disable-next-line @typescript-eslint/unbound-method
const setParams = useStore(state, (s) => s.setTxt2Img);
// eslint-disable-next-line @typescript-eslint/unbound-method
@ -50,22 +50,22 @@ export function Txt2Img() {
return <Box>
<Stack spacing={2}>
<Profiles
params={txt2img}
selectParams={selectParams}
selectHighres={selectHighres}
selectUpscale={selectUpscale}
setParams={setParams}
highres={highres}
setHighres={setHighres}
upscale={upscale}
setUpscale={setUpscale}
/>
<ModelControl model={model} setModel={setModel} />
<ImageControl selector={(s) => s.txt2img} onChange={setParams} />
<ImageControl selector={selectParams} onChange={setParams} />
<Stack direction='row' spacing={4}>
<NumericField
label={t('parameter.width')}
min={params.width.min}
max={params.width.max}
step={params.width.step}
value={txt2img.width}
value={width}
onChange={(value) => {
setParams({
width: value,
@ -77,7 +77,7 @@ export function Txt2Img() {
min={params.height.min}
max={params.height.max}
step={params.height.step}
value={txt2img.height}
value={height}
onChange={(value) => {
setParams({
height: value,
@ -85,8 +85,8 @@ export function Txt2Img() {
}}
/>
</Stack>
<HighresControl highres={highres} setHighres={setHighres} />
<UpscaleControl upscale={upscale} setUpscale={setUpscale} />
<HighresControl selectHighres={selectHighres} setHighres={setHighres} />
<UpscaleControl selectUpscale={selectUpscale} setUpscale={setUpscale} />
<Button
variant='contained'
onClick={() => generate.mutate()}

34
gui/src/components/tab/Upscale.tsx
View File

@ -6,8 +6,9 @@ import { useContext } from 'react';
import { useTranslation } from 'react-i18next';
import { useStore } from 'zustand';
import { HighresParams, ModelParams, UpscaleParams, UpscaleReqParams } from '../../client/types.js';
import { IMAGE_FILTER } from '../../config.js';
import { ClientContext, StateContext } from '../../state.js';
import { ClientContext, OnnxState, StateContext, TabState } from '../../state.js';
import { HighresControl } from '../control/HighresControl.js';
import { ModelControl } from '../control/ModelControl.js';
import { UpscaleControl } from '../control/UpscaleControl.js';
@ -33,10 +34,8 @@ export function Upscale() {
});
const state = mustExist(useContext(StateContext));
const highres = useStore(state, (s) => s.upscaleHighres);
const model = useStore(state, (s) => s.upscaleModel);
const params = useStore(state, (s) => s.upscale);
const upscale = useStore(state, (s) => s.upscaleUpscale);
// eslint-disable-next-line @typescript-eslint/unbound-method
const setModel = useStore(state, (s) => s.setUpscalingModel);
// eslint-disable-next-line @typescript-eslint/unbound-method
@ -51,7 +50,14 @@ export function Upscale() {
return <Box>
<Stack spacing={2}>
<Profiles params={params} setParams={setParams} highres={highres} setHighres={setHighres} upscale={upscale} setUpscale={setUpscale} />
<Profiles
selectHighres={selectHighres}
selectParams={selectParams}
selectUpscale={selectUpscale}
setParams={setParams}
setHighres={setHighres}
setUpscale={setUpscale}
/>
<ModelControl model={model} setModel={setModel} />
<ImageInput
filter={IMAGE_FILTER}
@ -70,8 +76,8 @@ export function Upscale() {
setParams(value);
}}
/>
<HighresControl highres={highres} setHighres={setHighres} />
<UpscaleControl upscale={upscale} setUpscale={setUpscale} />
<HighresControl selectHighres={selectHighres} setHighres={setHighres} />
<UpscaleControl selectUpscale={selectUpscale} setUpscale={setUpscale} />
<Button
disabled={doesExist(params.source) === false}
variant='contained'
@ -80,3 +86,19 @@ export function Upscale() {
</Stack>
</Box>;
}
export function selectModel(state: OnnxState): ModelParams {
return state.upscaleModel;
}
export function selectParams(state: OnnxState): TabState<UpscaleReqParams> {
return state.upscale;
}
export function selectHighres(state: OnnxState): HighresParams {
return state.upscaleHighres;
}
export function selectUpscale(state: OnnxState): UpscaleParams {
return state.upscaleUpscale;
}