From 1541818c1977a0ce4fbcc70ff7f0386dec19528b Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Thu, 21 Dec 2023 09:12:47 -0600
Subject: [PATCH] fix(gui): reduce renders when changing prompt and size params

---
 gui/src/components/control/ImageControl.tsx | 22 +------
 gui/src/components/input/PromptInput.tsx    | 64 +++++++++++---------
 gui/src/components/tab/Txt2Img.tsx          | 67 ++++++++++++---------
 3 files changed, 76 insertions(+), 77 deletions(-)

diff --git a/gui/src/components/control/ImageControl.tsx b/gui/src/components/control/ImageControl.tsx
index 8877ae50..9c8b2368 100644
--- a/gui/src/components/control/ImageControl.tsx
+++ b/gui/src/components/control/ImageControl.tsx
@@ -19,7 +19,7 @@ import { QueryList } from '../input/QueryList.js';
 
 const { useMemo } = React;
 
-type BaseParamsWithoutPrompt = Omit<BaseImgParams, 'prompt' | 'negativePrompt'>;
+type BaseParamsWithoutPrompt = Omit<BaseImgParams, 'prompt' | 'negativePrompt' | 'width' | 'height'>;
 
 export interface ImageControlProps {
   onChange(params: Partial<BaseImgParams>): void;
@@ -27,7 +27,7 @@ export interface ImageControlProps {
 }
 
 export function omitPrompt(selector: (state: OnnxState) => BaseImgParams): (state: OnnxState) => BaseParamsWithoutPrompt {
-  return (state) => omit(selector(state), 'prompt', 'negativePrompt');
+  return (state) => omit(selector(state), 'prompt', 'negativePrompt', 'width', 'height');
 }
 
 /**
@@ -77,7 +77,6 @@ export function ImageControl(props: ImageControlProps) {
         onChange={(eta) => {
           if (doesExist(onChange)) {
             onChange({
-              ...state,
               eta,
             });
           }
@@ -93,7 +92,6 @@ export function ImageControl(props: ImageControlProps) {
         onChange={(cfg) => {
           if (doesExist(onChange)) {
             onChange({
-              ...state,
               cfg,
             });
           }
@@ -107,7 +105,6 @@ export function ImageControl(props: ImageControlProps) {
         value={state.steps}
         onChange={(steps) => {
           onChange({
-            ...state,
             steps,
           });
         }}
@@ -120,7 +117,6 @@ export function ImageControl(props: ImageControlProps) {
         value={state.seed}
         onChange={(seed) => {
           onChange({
-            ...state,
             seed,
           });
         }}
@@ -131,7 +127,6 @@ export function ImageControl(props: ImageControlProps) {
         onClick={() => {
           const seed = Math.floor(Math.random() * params.seed.max);
           props.onChange({
-            ...state,
             seed,
           });
         }}
@@ -148,7 +143,6 @@ export function ImageControl(props: ImageControlProps) {
         value={state.batch}
         onChange={(batch) => {
           props.onChange({
-            ...state,
             batch,
           });
         }}
@@ -161,7 +155,6 @@ export function ImageControl(props: ImageControlProps) {
         value={state.unet_tile}
         onChange={(unet_tile) => {
           props.onChange({
-            ...state,
             unet_tile,
           });
         }}
@@ -175,7 +168,6 @@ export function ImageControl(props: ImageControlProps) {
         value={state.unet_overlap}
         onChange={(unet_overlap) => {
           props.onChange({
-            ...state,
             unet_overlap,
           });
         }}
@@ -187,7 +179,6 @@ export function ImageControl(props: ImageControlProps) {
           value='check'
           onChange={(event) => {
             props.onChange({
-              ...state,
               tiled_vae: state.tiled_vae === false,
             });
           }}
@@ -202,7 +193,6 @@ export function ImageControl(props: ImageControlProps) {
         value={state.vae_tile}
         onChange={(vae_tile) => {
           props.onChange({
-            ...state,
             vae_tile,
           });
         }}
@@ -217,7 +207,6 @@ export function ImageControl(props: ImageControlProps) {
         value={state.vae_overlap}
         onChange={(vae_overlap) => {
           props.onChange({
-            ...state,
             vae_overlap,
           });
         }}
@@ -225,12 +214,7 @@ export function ImageControl(props: ImageControlProps) {
     </Stack>
     <PromptInput
       selector={selector}
-      onChange={(value) => {
-        props.onChange({
-          ...state,
-          ...value,
-        });
-      }}
+      onChange={onChange}
     />
   </Stack>;
 }
diff --git a/gui/src/components/input/PromptInput.tsx b/gui/src/components/input/PromptInput.tsx
index b32d4474..9661af1b 100644
--- a/gui/src/components/input/PromptInput.tsx
+++ b/gui/src/components/input/PromptInput.tsx
@@ -7,10 +7,11 @@ import { useTranslation } from 'react-i18next';
 import { useStore } from 'zustand';
 import { shallow } from 'zustand/shallow';
 
+import { memo, useCallback } from 'react';
 import { STALE_TIME } from '../../config.js';
 import { ClientContext, OnnxState, StateContext } from '../../state/full.js';
-import { QueryMenu } from '../input/QueryMenu.js';
 import { ModelResponse, NetworkModel } from '../../types/api.js';
+import { QueryMenu, QueryMenuComplete, QueryMenuFilter } from '../input/QueryMenu.js';
 
 const { useContext, useMemo } = React;
 
@@ -89,6 +90,9 @@ export function PromptTextBlock(props: PromptTextBlockProps) {
   </Stack>;
 }
 
+const ModelMenu = memo(QueryMenu<ModelResponse>);
+const StringMenu = memo(QueryMenu<Array<string>>);
+
 export function PromptInput(props: PromptInputProps) {
   // eslint-disable-next-line @typescript-eslint/unbound-method
   const { selector, onChange } = props;
@@ -104,21 +108,38 @@ export function PromptInput(props: PromptInputProps) {
 
   const { t } = useTranslation();
 
-  function addNetwork(type: string, name: string, weight = 1.0) {
+  const addNetwork = useCallback((type: string, name: string, weight = 1.0) => {
     const { prompt, negativePrompt } = selector(store.getState());
     onChange({
       negativePrompt,
       prompt: `<${type}:${name}:${weight.toFixed(2)}> ${prompt}`,
     });
-  }
+  }, [ onChange ]);
 
-  function addWildcard(name: string) {
+  const addInversion = useCallback((name: string) => addNetwork('inversion', name), [ onChange ]);
+  const addLora = useCallback((name: string) => addNetwork('lora', name), [ onChange ]);
+
+  const addWildcard = useCallback((name: string) => {
     const { prompt, negativePrompt } = selector(store.getState());
     onChange({
       negativePrompt,
       prompt: `${prompt}, __${name}__`,
     });
-  }
+  }, [ onChange ]);
+
+  const inversionSelector = useMemo<QueryMenuFilter<ModelResponse>>(() => ({
+    result: models,
+    selector: (result) => filterNetworks(result.networks, 'inversion'),
+  }), [models.status]);
+
+  const loraSelector = useMemo<QueryMenuFilter<ModelResponse>>(() => ({
+    result: models,
+    selector: (result) => filterNetworks(result.networks, 'lora'),
+  }), [models.status]);
+
+  const wildcardSelector = useMemo<QueryMenuComplete>(() => ({
+    result: wildcards,
+  }), [wildcards.status]);
 
   return <Stack spacing={2}>
     <PromptTextBlock
@@ -127,41 +148,26 @@ export function PromptInput(props: PromptInputProps) {
       selector={selector}
     />
     <Stack direction='row' spacing={2}>
-      <QueryMenu
+      <ModelMenu
         id='inversion'
         labelKey='model.inversion'
         name={t('modelType.inversion')}
-        query={{
-          result: models,
-          selector: (result) => filterNetworks(result.networks, 'inversion'),
-        }}
-        onSelect={(name) => {
-          addNetwork('inversion', name);
-        }}
+        query={inversionSelector}
+        onSelect={addInversion}
       />
-      <QueryMenu
+      <ModelMenu
         id='lora'
         labelKey='model.lora'
         name={t('modelType.lora')}
-        query={{
-          result: models,
-          selector: (result) => filterNetworks(result.networks, 'lora'),
-        }}
-        onSelect={(name) => {
-          addNetwork('lora', name);
-        }}
+        query={loraSelector}
+        onSelect={addLora}
       />
-      <QueryMenu
+      <StringMenu
         id='wildcard'
         labelKey='wildcard'
         name={t('wildcard')}
-        query={{
-          result: wildcards,
-          selector: (result) => result,
-        }}
-        onSelect={(name) => {
-          addWildcard(name);
-        }}
+        query={wildcardSelector}
+        onSelect={addWildcard}
       />
     </Stack>
   </Stack>;
diff --git a/gui/src/components/tab/Txt2Img.tsx b/gui/src/components/tab/Txt2Img.tsx
index bcae3576..210cbaa3 100644
--- a/gui/src/components/tab/Txt2Img.tsx
+++ b/gui/src/components/tab/Txt2Img.tsx
@@ -19,9 +19,44 @@ import { UpscaleControl } from '../control/UpscaleControl.js';
 import { VariableControl } from '../control/VariableControl.js';
 import { NumericField } from '../input/NumericField.js';
 
-export function Txt2Img() {
+export function SizeControl() {
   const { params } = mustExist(useContext(ConfigContext));
 
+  const store = mustExist(useContext(StateContext));
+  const { height, width } = useStore(store, selectSize, shallow);
+  const { setParams } = useStore(store, selectActions, shallow);
+
+  const { t } = useTranslation();
+
+  return <Stack direction='row' spacing={4}>
+    <NumericField
+      label={t('parameter.width')}
+      min={params.width.min}
+      max={params.width.max}
+      step={params.width.step}
+      value={width}
+      onChange={(value) => {
+        setParams({
+          width: value,
+        });
+      }}
+    />
+    <NumericField
+      label={t('parameter.height')}
+      min={params.height.min}
+      max={params.height.max}
+      step={params.height.step}
+      value={height}
+      onChange={(value) => {
+        setParams({
+          height: value,
+        });
+      }}
+    />
+  </Stack>;
+}
+
+export function Txt2Img() {
   async function generateImage() {
     const state = store.getState();
     const grid = selectVariable(state);
@@ -47,7 +82,6 @@ export function Txt2Img() {
 
   const store = mustExist(useContext(StateContext));
   const { pushHistory, setHighres, setModel, setParams, setUpscale, setVariable } = useStore(store, selectActions, shallow);
-  const { height, width } = useStore(store, selectReactParams, shallow);
   const model = useStore(store, selectModel);
 
   const { t } = useTranslation();
@@ -66,32 +100,7 @@ export function Txt2Img() {
       />
       <ModelControl model={model} setModel={setModel} />
       <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={width}
-          onChange={(value) => {
-            setParams({
-              width: value,
-            });
-          }}
-        />
-        <NumericField
-          label={t('parameter.height')}
-          min={params.height.min}
-          max={params.height.max}
-          step={params.height.step}
-          value={height}
-          onChange={(value) => {
-            setParams({
-              height: value,
-            });
-          }}
-        />
-      </Stack>
+      <SizeControl />
       <HighresControl selectHighres={selectHighres} setHighres={setHighres} />
       <UpscaleControl selectUpscale={selectUpscale} setUpscale={setUpscale} />
       <VariableControl selectGrid={selectVariable} setGrid={setVariable} />
@@ -128,7 +137,7 @@ export function selectParams(state: OnnxState): TabState<Txt2ImgParams> {
   return state.txt2img;
 }
 
-export function selectReactParams(state: OnnxState) {
+export function selectSize(state: OnnxState) {
   return {
     height: state.txt2img.height,
     width: state.txt2img.width,