From 83ecda977ca4a66caf098818856f7a1a41608a3a Mon Sep 17 00:00:00 2001
From: blessedcoolant <54517381+blessedcoolant@users.noreply.github.com>
Date: Thu, 9 Feb 2023 20:19:25 +1300
Subject: [PATCH] Add frontend UI for denoise_str for ESRGAN

---
 .../public/locales/parameters/en.json         |  1 +
 .../frontend/src/app/socketio/emitters.ts     |  8 ++-
 .../src/common/util/parameterTranslation.ts   |  7 +--
 .../Upscale/UpscaleSettings.scss              |  5 --
 .../Upscale/UpscaleSettings.tsx               | 58 ++++++++++++++-----
 .../parameters/store/postprocessingSlice.ts   |  6 ++
 invokeai/frontend/src/styles/index.scss       |  1 -
 7 files changed, 61 insertions(+), 25 deletions(-)
 delete mode 100644 invokeai/frontend/src/features/parameters/components/AdvancedParameters/Upscale/UpscaleSettings.scss

diff --git a/invokeai/frontend/public/locales/parameters/en.json b/invokeai/frontend/public/locales/parameters/en.json
index 5bca8b0950..fef06c92fa 100644
--- a/invokeai/frontend/public/locales/parameters/en.json
+++ b/invokeai/frontend/public/locales/parameters/en.json
@@ -20,6 +20,7 @@
   "upscaling": "Upscaling",
   "upscale": "Upscale",
   "upscaleImage": "Upscale Image",
+  "denoisingStrength": "Denoising Strength",
   "scale": "Scale",
   "otherOptions": "Other Options",
   "seamlessTiling": "Seamless Tiling",
diff --git a/invokeai/frontend/src/app/socketio/emitters.ts b/invokeai/frontend/src/app/socketio/emitters.ts
index 6978fa9d2b..bc90f9c58b 100644
--- a/invokeai/frontend/src/app/socketio/emitters.ts
+++ b/invokeai/frontend/src/app/socketio/emitters.ts
@@ -93,11 +93,15 @@ const makeSocketIOEmitters = (
       dispatch(setIsProcessing(true));
 
       const {
-        postprocessing: { upscalingLevel, upscalingStrength },
+        postprocessing: {
+          upscalingLevel,
+          upscalingDenoising,
+          upscalingStrength,
+        },
       } = getState();
 
       const esrganParameters = {
-        upscale: [upscalingLevel, upscalingStrength],
+        upscale: [upscalingLevel, upscalingDenoising, upscalingStrength],
       };
       socketio.emit('runPostprocessing', imageToProcess, {
         type: 'esrgan',
diff --git a/invokeai/frontend/src/common/util/parameterTranslation.ts b/invokeai/frontend/src/common/util/parameterTranslation.ts
index 1e429d8958..60d3c1e485 100644
--- a/invokeai/frontend/src/common/util/parameterTranslation.ts
+++ b/invokeai/frontend/src/common/util/parameterTranslation.ts
@@ -69,6 +69,7 @@ export type BackendGenerationParameters = {
 
 export type BackendEsrGanParameters = {
   level: UpscalingLevel;
+  denoise_str: number;
   strength: number;
 };
 
@@ -111,13 +112,12 @@ export const frontendToBackendParameters = (
     shouldRunFacetool,
     upscalingLevel,
     upscalingStrength,
+    upscalingDenoising,
   } = postprocessingState;
 
   const {
     cfgScale,
-
     height,
-
     img2imgStrength,
     infillMethod,
     initialImage,
@@ -136,11 +136,9 @@ export const frontendToBackendParameters = (
     shouldFitToWidthHeight,
     shouldGenerateVariations,
     shouldRandomizeSeed,
-
     steps,
     threshold,
     tileSize,
-
     variationAmount,
     width,
   } = generationState;
@@ -190,6 +188,7 @@ export const frontendToBackendParameters = (
     if (shouldRunESRGAN) {
       esrganParameters = {
         level: upscalingLevel,
+        denoise_str: upscalingDenoising,
         strength: upscalingStrength,
       };
     }
diff --git a/invokeai/frontend/src/features/parameters/components/AdvancedParameters/Upscale/UpscaleSettings.scss b/invokeai/frontend/src/features/parameters/components/AdvancedParameters/Upscale/UpscaleSettings.scss
deleted file mode 100644
index 921253e5ee..0000000000
--- a/invokeai/frontend/src/features/parameters/components/AdvancedParameters/Upscale/UpscaleSettings.scss
+++ /dev/null
@@ -1,5 +0,0 @@
-.upscale-settings {
-  display: grid;
-  grid-template-columns: auto 1fr;
-  column-gap: 1rem;
-}
diff --git a/invokeai/frontend/src/features/parameters/components/AdvancedParameters/Upscale/UpscaleSettings.tsx b/invokeai/frontend/src/features/parameters/components/AdvancedParameters/Upscale/UpscaleSettings.tsx
index 180941597c..bb987d7eb5 100644
--- a/invokeai/frontend/src/features/parameters/components/AdvancedParameters/Upscale/UpscaleSettings.tsx
+++ b/invokeai/frontend/src/features/parameters/components/AdvancedParameters/Upscale/UpscaleSettings.tsx
@@ -1,6 +1,7 @@
 import { useAppDispatch, useAppSelector } from 'app/storeHooks';
 
 import {
+  setUpscalingDenoising,
   setUpscalingLevel,
   setUpscalingStrength,
   UpscalingLevel,
@@ -8,20 +9,25 @@ import {
 
 import { createSelector } from '@reduxjs/toolkit';
 import { UPSCALING_LEVELS } from 'app/constants';
-import IAINumberInput from 'common/components/IAINumberInput';
 import IAISelect from 'common/components/IAISelect';
 import { postprocessingSelector } from 'features/parameters/store/postprocessingSelectors';
 import { systemSelector } from 'features/system/store/systemSelectors';
 import { isEqual } from 'lodash';
 import { ChangeEvent } from 'react';
 import { useTranslation } from 'react-i18next';
+import IAISlider from 'common/components/IAISlider';
+import { Flex } from '@chakra-ui/react';
 
 const parametersSelector = createSelector(
   [postprocessingSelector, systemSelector],
 
-  ({ upscalingLevel, upscalingStrength }, { isESRGANAvailable }) => {
+  (
+    { upscalingLevel, upscalingStrength, upscalingDenoising },
+    { isESRGANAvailable }
+  ) => {
     return {
       upscalingLevel,
+      upscalingDenoising,
       upscalingStrength,
       isESRGANAvailable,
     };
@@ -38,8 +44,12 @@ const parametersSelector = createSelector(
  */
 const UpscaleSettings = () => {
   const dispatch = useAppDispatch();
-  const { upscalingLevel, upscalingStrength, isESRGANAvailable } =
-    useAppSelector(parametersSelector);
+  const {
+    upscalingLevel,
+    upscalingStrength,
+    upscalingDenoising,
+    isESRGANAvailable,
+  } = useAppSelector(parametersSelector);
 
   const { t } = useTranslation();
 
@@ -49,7 +59,7 @@ const UpscaleSettings = () => {
   const handleChangeStrength = (v: number) => dispatch(setUpscalingStrength(v));
 
   return (
-    <div className="upscale-settings">
+    <Flex flexDir="column" rowGap="1rem" minWidth="20rem">
       <IAISelect
         isDisabled={!isESRGANAvailable}
         label={t('parameters:scale')}
@@ -57,17 +67,39 @@ const UpscaleSettings = () => {
         onChange={handleChangeLevel}
         validValues={UPSCALING_LEVELS}
       />
-      <IAINumberInput
-        isDisabled={!isESRGANAvailable}
-        label={t('parameters:strength')}
-        step={0.05}
+      <IAISlider
+        label={t('parameters:denoisingStrength')}
+        value={upscalingDenoising}
         min={0}
         max={1}
-        onChange={handleChangeStrength}
-        value={upscalingStrength}
-        isInteger={false}
+        step={0.01}
+        onChange={(v) => {
+          dispatch(setUpscalingDenoising(v));
+        }}
+        handleReset={() => dispatch(setUpscalingDenoising(0.75))}
+        withSliderMarks
+        withInput
+        withReset
+        isSliderDisabled={!isESRGANAvailable}
+        isInputDisabled={!isESRGANAvailable}
+        isResetDisabled={!isESRGANAvailable}
       />
-    </div>
+      <IAISlider
+        label={`${t('parameters:upscale')} ${t('parameters:strength')}`}
+        value={upscalingStrength}
+        min={0}
+        max={1}
+        step={0.05}
+        onChange={handleChangeStrength}
+        handleReset={() => dispatch(setUpscalingStrength(0.75))}
+        withSliderMarks
+        withInput
+        withReset
+        isSliderDisabled={!isESRGANAvailable}
+        isInputDisabled={!isESRGANAvailable}
+        isResetDisabled={!isESRGANAvailable}
+      />
+    </Flex>
   );
 };
 
diff --git a/invokeai/frontend/src/features/parameters/store/postprocessingSlice.ts b/invokeai/frontend/src/features/parameters/store/postprocessingSlice.ts
index 9c7290ec3f..7f9e0c3aa2 100644
--- a/invokeai/frontend/src/features/parameters/store/postprocessingSlice.ts
+++ b/invokeai/frontend/src/features/parameters/store/postprocessingSlice.ts
@@ -16,6 +16,7 @@ export interface PostprocessingState {
   shouldRunESRGAN: boolean;
   shouldRunFacetool: boolean;
   upscalingLevel: UpscalingLevel;
+  upscalingDenoising: number;
   upscalingStrength: number;
 }
 
@@ -29,6 +30,7 @@ const initialPostprocessingState: PostprocessingState = {
   shouldRunESRGAN: false,
   shouldRunFacetool: false,
   upscalingLevel: 4,
+  upscalingDenoising: 0.75,
   upscalingStrength: 0.75,
 };
 
@@ -47,6 +49,9 @@ export const postprocessingSlice = createSlice({
     setUpscalingLevel: (state, action: PayloadAction<UpscalingLevel>) => {
       state.upscalingLevel = action.payload;
     },
+    setUpscalingDenoising: (state, action: PayloadAction<number>) => {
+      state.upscalingDenoising = action.payload;
+    },
     setUpscalingStrength: (state, action: PayloadAction<number>) => {
       state.upscalingStrength = action.payload;
     },
@@ -88,6 +93,7 @@ export const {
   setShouldRunESRGAN,
   setShouldRunFacetool,
   setUpscalingLevel,
+  setUpscalingDenoising,
   setUpscalingStrength,
 } = postprocessingSlice.actions;
 
diff --git a/invokeai/frontend/src/styles/index.scss b/invokeai/frontend/src/styles/index.scss
index d93359bf2a..afafc17181 100644
--- a/invokeai/frontend/src/styles/index.scss
+++ b/invokeai/frontend/src/styles/index.scss
@@ -27,7 +27,6 @@
 @use '../features/parameters/components/ProcessButtons/ProcessButtons.scss';
 @use '../features/parameters/components/MainParameters/MainParameters.scss';
 @use '../features/parameters/components/AccordionItems/AdvancedSettings.scss';
-@use '../features/parameters/components/AdvancedParameters/Upscale/UpscaleSettings.scss';
 @use '../features/parameters/components/AdvancedParameters/Canvas/BoundingBox/BoundingBoxSettings.scss';
 
 // gallery