optimize zustand selectors

gui/src/components/Profiles.tsx

@@ -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;
+  selectHighres(state: OnnxState): HighresParams;
-  params: BaseImgParams;
+  selectParams(state: OnnxState): BaseImgParams;
-  upscale: UpscaleParams;
+  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,
+              highres: props.selectHighres(innerState),
-              upscale: props.upscale,
+              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>

gui/src/components/control/HighresControl.tsx

@@ -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();

gui/src/components/control/UpscaleControl.tsx

@@ -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();

gui/src/components/tab/Blend.tsx

@@ -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 blend = useStore(state, selectParams);
-  const upscale = useStore(state, (s) => s.blendUpscale);
   // 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;
+}

gui/src/components/tab/Img2Img.tsx

@@ -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 sourceFilter = useStore(state, (s) => s.img2img.sourceFilter);
-  const highres = useStore(state, selectHighres);
+  const strength = useStore(state, (s) => s.img2img.strength);
-  const upscale = useStore(state, selectUpscale);
+  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} />
+      <HighresControl selectHighres={selectHighres} setHighres={setHighres} />
-      <UpscaleControl upscale={upscale} setUpscale={setUpscale} />
+      <UpscaleControl selectUpscale={selectUpscale} setUpscale={setUpscale} />
       <Button
         disabled={doesExist(source) === false}
         variant='contained'

gui/src/components/tab/Inpaint.tsx

@@ -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),
+        mask: mustExist(mask),
-        source: mustExist(inpaint.source),
+        source: mustExist(source),
-      }, upscale, highres);
+      }, selectUpscale(innerState), selectHighres(innerState));
 
       pushHistory(image, retry);
     } else {
       const { image, retry } = await client.inpaint(model, {
         ...inpaint,
-        mask: mustExist(inpaint.mask),
+        mask: mustExist(mask),
-        source: mustExist(inpaint.source),
+        source: mustExist(source),
-      }, upscale, highres);
+      }, 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 source = useStore(state, (s) => s.inpaint.source);
-  const highres = useStore(state, selectHighres);
+  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}
+        source={source}
-        mask={inpaint.mask}
+        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} />
+      <HighresControl selectHighres={selectHighres} setHighres={setHighres} />
-      <UpscaleControl upscale={upscale} setUpscale={setUpscale} />
+      <UpscaleControl selectUpscale={selectUpscale} setUpscale={setUpscale} />
       <Button
         disabled={preventInpaint()}
         variant='contained'

gui/src/components/tab/Txt2Img.tsx

@@ -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} />
+      <HighresControl selectHighres={selectHighres} setHighres={setHighres} />
-      <UpscaleControl upscale={upscale} setUpscale={setUpscale} />
+      <UpscaleControl selectUpscale={selectUpscale} setUpscale={setUpscale} />
       <Button
         variant='contained'
         onClick={() => generate.mutate()}

gui/src/components/tab/Upscale.tsx

@@ -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} />
+      <HighresControl selectHighres={selectHighres} setHighres={setHighres} />
-      <UpscaleControl upscale={upscale} setUpscale={setUpscale} />
+      <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;
+}