From d3ad43bef4914df77a95e65ff353e3f7862df6e3 Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Sun, 15 Jan 2023 09:34:17 -0600
Subject: [PATCH] feat(gui): add menus for noise source and blend mode

---
 gui/src/api/client.ts          | 19 ++++++++++++++++++
 gui/src/components/Inpaint.tsx | 36 ++++++++++++++++++++++++++++++++--
 gui/src/state.ts               |  4 ++++
 gui/src/strings.ts             | 13 ++++++++++++
 4 files changed, 70 insertions(+), 2 deletions(-)

diff --git a/gui/src/api/client.ts b/gui/src/api/client.ts
index 35f088a7..f7ef3013 100644
--- a/gui/src/api/client.ts
+++ b/gui/src/api/client.ts
@@ -43,6 +43,9 @@ export type Txt2ImgResponse = Required<Txt2ImgParams>;
 export interface InpaintParams extends BaseImgParams {
   mask: Blob;
   source: Blob;
+
+  blend: string;
+  noise: string;
 }
 
 export interface OutpaintPixels {
@@ -75,7 +78,9 @@ export interface ApiReady {
 }
 
 export interface ApiClient {
+  blends(): Promise<Array<string>>;
   models(): Promise<Array<string>>;
+  noises(): Promise<Array<string>>;
   params(): Promise<ConfigParams>;
   platforms(): Promise<Array<string>>;
   schedulers(): Promise<Array<string>>;
@@ -156,11 +161,21 @@ export function makeClient(root: string, f = fetch): ApiClient {
   }
 
   return {
+    async blends(): Promise<Array<string>> {
+      const path = makeApiUrl(root, 'settings', 'blends');
+      const res = await f(path);
+      return await res.json() as Array<string>;
+    },
     async models(): Promise<Array<string>> {
       const path = makeApiUrl(root, 'settings', 'models');
       const res = await f(path);
       return await res.json() as Array<string>;
     },
+    async noises(): Promise<Array<string>> {
+      const path = makeApiUrl(root, 'settings', 'noises');
+      const res = await f(path);
+      return await res.json() as Array<string>;
+    },
     async params(): Promise<ConfigParams> {
       const path = makeApiUrl(root, 'settings', 'params');
       const res = await f(path);
@@ -223,6 +238,8 @@ export function makeClient(root: string, f = fetch): ApiClient {
       }
 
       const url = makeImageURL(root, 'inpaint', params);
+      url.searchParams.append('noise', params.noise);
+      url.searchParams.append('blend', params.blend);
 
       const body = new FormData();
       body.append('mask', params.mask, 'mask');
@@ -242,6 +259,8 @@ export function makeClient(root: string, f = fetch): ApiClient {
       }
 
       const url = makeImageURL(root, 'inpaint', params);
+      url.searchParams.append('noise', params.noise);
+      url.searchParams.append('blend', params.blend);
 
       if (doesExist(params.left)) {
         url.searchParams.append('left', params.left.toFixed(0));
diff --git a/gui/src/components/Inpaint.tsx b/gui/src/components/Inpaint.tsx
index c235e7cf..dd039f8e 100644
--- a/gui/src/components/Inpaint.tsx
+++ b/gui/src/components/Inpaint.tsx
@@ -1,15 +1,17 @@
 import { mustExist } from '@apextoaster/js-utils';
 import { Box, Button, Stack } from '@mui/material';
 import * as React from 'react';
-import { useMutation, useQueryClient } from 'react-query';
+import { useMutation, useQuery, useQueryClient } from 'react-query';
 import { useStore } from 'zustand';
 
-import { ConfigParams, IMAGE_FILTER } from '../config.js';
+import { ConfigParams, IMAGE_FILTER, STALE_TIME } from '../config.js';
 import { ClientContext, StateContext } from '../state.js';
+import { BLEND_LABELS, NOISE_LABELS } from '../strings.js';
 import { ImageControl } from './ImageControl.js';
 import { ImageInput } from './ImageInput.js';
 import { MaskCanvas } from './MaskCanvas.js';
 import { OutpaintControl } from './OutpaintControl.js';
+import { QueryList } from './QueryList.js';
 
 const { useContext } = React;
 
@@ -23,6 +25,12 @@ export interface InpaintProps {
 export function Inpaint(props: InpaintProps) {
   const { config, model, platform } = props;
   const client = mustExist(useContext(ClientContext));
+  const blends = useQuery('blends', async () => client.blends(), {
+    staleTime: STALE_TIME,
+  });
+  const noises = useQuery('noises', async () => client.noises(), {
+    staleTime: STALE_TIME,
+  });
 
   async function uploadSource(): Promise<void> {
     const outpaint = state.getState().outpaint; // TODO: seems shady
@@ -104,6 +112,30 @@ export function Inpaint(props: InpaintProps) {
           setInpaint(newParams);
         }}
       />
+      <QueryList
+        id='blends'
+        labels={BLEND_LABELS}
+        name='Blend Mode'
+        result={blends}
+        value={params.blend}
+        onChange={(blend) => {
+          setInpaint({
+            blend,
+          });
+        }}
+      />
+      <QueryList
+        id='noises'
+        labels={NOISE_LABELS}
+        name='Noise Source'
+        result={noises}
+        value={params.noise}
+        onChange={(noise) => {
+          setInpaint({
+            noise,
+          });
+        }}
+      />
       <OutpaintControl config={config} />
       <Button onClick={() => upload.mutate()}>Generate</Button>
     </Stack>
diff --git a/gui/src/state.ts b/gui/src/state.ts
index edf8e487..d7726b98 100644
--- a/gui/src/state.ts
+++ b/gui/src/state.ts
@@ -128,6 +128,8 @@ export function createStateSlices(base: ConfigParams) {
       ...defaults,
       mask: null,
       source: null,
+      noise: '',
+      blend: '',
     },
     setInpaint(params) {
       set((prev) => ({
@@ -143,6 +145,8 @@ export function createStateSlices(base: ConfigParams) {
           ...defaults,
           mask: null,
           source: null,
+          noise: '',
+          blend: '',
         },
       });
     },
diff --git a/gui/src/strings.ts b/gui/src/strings.ts
index 6d12594a..4d49a492 100644
--- a/gui/src/strings.ts
+++ b/gui/src/strings.ts
@@ -26,3 +26,16 @@ export const SCHEDULER_LABELS: Record<string, string> = {
   'lms-discrete': 'LMS',
   'pndm': 'PNDM',
 };
+
+export const NOISE_LABELS: Record<string, string> = {
+  fill: 'Fill Edges',
+  guassian: 'Gaussian Blur',
+  histogram: 'Histogram Noise',
+  normal: 'Gaussian Noise',
+  uniform: 'Uniform Noise',
+};
+
+export const BLEND_LABELS: Record<string, string> = {
+  'mask-source': 'Noise -> Source',
+  'source-mask': 'Source -> Noise',
+};