From 497f66e6823e1adf119cb46590df8cadd810e837 Mon Sep 17 00:00:00 2001
From: blessedcoolant <54517381+blessedcoolant@users.noreply.github.com>
Date: Sat, 2 Sep 2023 10:24:32 +1200
Subject: [PATCH] feat: Add Patchmatch Downscale control to UI + refine the ui
 there

---
 invokeai/frontend/web/public/locales/en.json  |  1 +
 .../graphBuilders/buildCanvasOutpaintGraph.ts |  6 +-
 .../buildCanvasSDXLOutpaintGraph.ts           |  5 +-
 .../ParamInfillAndScalingCollapse.tsx         |  4 +-
 .../InfillAndScaling/ParamInfillOptions.tsx   | 29 +++++++++
 .../ParamInfillPatchmatchDownscaleSize.tsx    | 59 +++++++++++++++++++
 .../InfillAndScaling/ParamInfillTilesize.tsx  | 14 ++---
 .../parameters/store/generationSlice.ts       | 21 +++++--
 8 files changed, 120 insertions(+), 19 deletions(-)
 create mode 100644 invokeai/frontend/web/src/features/parameters/components/Parameters/Canvas/InfillAndScaling/ParamInfillOptions.tsx
 create mode 100644 invokeai/frontend/web/src/features/parameters/components/Parameters/Canvas/InfillAndScaling/ParamInfillPatchmatchDownscaleSize.tsx

diff --git a/invokeai/frontend/web/public/locales/en.json b/invokeai/frontend/web/public/locales/en.json
index 32144749bd..bcca3fc8e7 100644
--- a/invokeai/frontend/web/public/locales/en.json
+++ b/invokeai/frontend/web/public/locales/en.json
@@ -520,6 +520,7 @@
         "scaledHeight": "Scaled H",
         "infillMethod": "Infill Method",
         "tileSize": "Tile Size",
+        "patchmatchDownScaleSize": "Downscale",
         "boundingBoxHeader": "Bounding Box",
         "seamCorrectionHeader": "Seam Correction",
         "infillScalingHeader": "Infill and Scaling",
diff --git a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildCanvasOutpaintGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildCanvasOutpaintGraph.ts
index 64c60206fd..20f81611fe 100644
--- a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildCanvasOutpaintGraph.ts
+++ b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildCanvasOutpaintGraph.ts
@@ -73,7 +73,8 @@ export const buildCanvasOutpaintGraph = (
     maskBlur,
     canvasCoherenceSteps,
     canvasCoherenceStrength,
-    tileSize,
+    infillTileSize,
+    infillPatchmatchDownscaleSize,
     infillMethod,
     clipSkip,
     seamlessXAxis,
@@ -495,6 +496,7 @@ export const buildCanvasOutpaintGraph = (
       type: 'infill_patchmatch',
       id: INPAINT_INFILL,
       is_intermediate: true,
+      downscale: infillPatchmatchDownscaleSize,
     };
   }
 
@@ -519,7 +521,7 @@ export const buildCanvasOutpaintGraph = (
       type: 'infill_tile',
       id: INPAINT_INFILL,
       is_intermediate: true,
-      tile_size: tileSize,
+      tile_size: infillTileSize,
     };
   }
 
diff --git a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildCanvasSDXLOutpaintGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildCanvasSDXLOutpaintGraph.ts
index 6ab1e7e7b2..bd06228beb 100644
--- a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildCanvasSDXLOutpaintGraph.ts
+++ b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildCanvasSDXLOutpaintGraph.ts
@@ -74,7 +74,7 @@ export const buildCanvasSDXLOutpaintGraph = (
     maskBlur,
     canvasCoherenceSteps,
     canvasCoherenceStrength,
-    tileSize,
+    infillTileSize,
     infillMethod,
     seamlessXAxis,
     seamlessYAxis,
@@ -510,6 +510,7 @@ export const buildCanvasSDXLOutpaintGraph = (
       type: 'infill_patchmatch',
       id: INPAINT_INFILL,
       is_intermediate: true,
+      downscale: infillPatchmatchDownscaleSize,
     };
   }
 
@@ -534,7 +535,7 @@ export const buildCanvasSDXLOutpaintGraph = (
       type: 'infill_tile',
       id: INPAINT_INFILL,
       is_intermediate: true,
-      tile_size: tileSize,
+      tile_size: infillTileSize,
     };
   }
 
diff --git a/invokeai/frontend/web/src/features/parameters/components/Parameters/Canvas/InfillAndScaling/ParamInfillAndScalingCollapse.tsx b/invokeai/frontend/web/src/features/parameters/components/Parameters/Canvas/InfillAndScaling/ParamInfillAndScalingCollapse.tsx
index 8907e46e5f..b9c5e218ce 100644
--- a/invokeai/frontend/web/src/features/parameters/components/Parameters/Canvas/InfillAndScaling/ParamInfillAndScalingCollapse.tsx
+++ b/invokeai/frontend/web/src/features/parameters/components/Parameters/Canvas/InfillAndScaling/ParamInfillAndScalingCollapse.tsx
@@ -5,7 +5,7 @@ import { useTranslation } from 'react-i18next';
 import IAICollapse from 'common/components/IAICollapse';
 import SubParametersWrapper from '../../SubParametersWrapper';
 import ParamInfillMethod from './ParamInfillMethod';
-import ParamInfillTilesize from './ParamInfillTilesize';
+import ParamInfillOptions from './ParamInfillOptions';
 import ParamScaleBeforeProcessing from './ParamScaleBeforeProcessing';
 import ParamScaledHeight from './ParamScaledHeight';
 import ParamScaledWidth from './ParamScaledWidth';
@@ -18,7 +18,7 @@ const ParamInfillCollapse = () => {
       <Flex sx={{ gap: 2, flexDirection: 'column' }}>
         <SubParametersWrapper>
           <ParamInfillMethod />
-          <ParamInfillTilesize />
+          <ParamInfillOptions />
         </SubParametersWrapper>
         <Divider />
         <SubParametersWrapper>
diff --git a/invokeai/frontend/web/src/features/parameters/components/Parameters/Canvas/InfillAndScaling/ParamInfillOptions.tsx b/invokeai/frontend/web/src/features/parameters/components/Parameters/Canvas/InfillAndScaling/ParamInfillOptions.tsx
new file mode 100644
index 0000000000..14c2663d5e
--- /dev/null
+++ b/invokeai/frontend/web/src/features/parameters/components/Parameters/Canvas/InfillAndScaling/ParamInfillOptions.tsx
@@ -0,0 +1,29 @@
+import { Flex } from '@chakra-ui/react';
+import { createSelector } from '@reduxjs/toolkit';
+import { useAppSelector } from 'app/store/storeHooks';
+import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
+import { generationSelector } from 'features/parameters/store/generationSelectors';
+import ParamInfillPatchmatchDownscaleSize from './ParamInfillPatchmatchDownscaleSize';
+import ParamInfillTilesize from './ParamInfillTilesize';
+
+const selector = createSelector(
+  [generationSelector],
+  (parameters) => {
+    const { infillMethod } = parameters;
+
+    return {
+      infillMethod,
+    };
+  },
+  defaultSelectorOptions
+);
+
+export default function ParamInfillOptions() {
+  const { infillMethod } = useAppSelector(selector);
+  return (
+    <Flex>
+      {infillMethod === 'tile' && <ParamInfillTilesize />}
+      {infillMethod === 'patchmatch' && <ParamInfillPatchmatchDownscaleSize />}
+    </Flex>
+  );
+}
diff --git a/invokeai/frontend/web/src/features/parameters/components/Parameters/Canvas/InfillAndScaling/ParamInfillPatchmatchDownscaleSize.tsx b/invokeai/frontend/web/src/features/parameters/components/Parameters/Canvas/InfillAndScaling/ParamInfillPatchmatchDownscaleSize.tsx
new file mode 100644
index 0000000000..d8e9749422
--- /dev/null
+++ b/invokeai/frontend/web/src/features/parameters/components/Parameters/Canvas/InfillAndScaling/ParamInfillPatchmatchDownscaleSize.tsx
@@ -0,0 +1,59 @@
+import { createSelector } from '@reduxjs/toolkit';
+import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
+import IAISlider from 'common/components/IAISlider';
+import { generationSelector } from 'features/parameters/store/generationSelectors';
+import { setInfillPatchmatchDownscaleSize } from 'features/parameters/store/generationSlice';
+import { memo, useCallback } from 'react';
+
+import { useTranslation } from 'react-i18next';
+
+const selector = createSelector(
+  [generationSelector],
+  (parameters) => {
+    const { infillPatchmatchDownscaleSize, infillMethod } = parameters;
+
+    return {
+      infillPatchmatchDownscaleSize,
+      infillMethod,
+    };
+  },
+  defaultSelectorOptions
+);
+
+const ParamInfillPatchmatchDownscaleSize = () => {
+  const dispatch = useAppDispatch();
+  const { infillPatchmatchDownscaleSize, infillMethod } =
+    useAppSelector(selector);
+
+  const { t } = useTranslation();
+
+  const handleChange = useCallback(
+    (v: number) => {
+      dispatch(setInfillPatchmatchDownscaleSize(v));
+    },
+    [dispatch]
+  );
+
+  const handleReset = useCallback(() => {
+    dispatch(setInfillPatchmatchDownscaleSize(2));
+  }, [dispatch]);
+
+  return (
+    <IAISlider
+      isDisabled={infillMethod !== 'patchmatch'}
+      label={t('parameters.patchmatchDownScaleSize')}
+      min={1}
+      max={10}
+      sliderNumberInputProps={{ max: 10 }}
+      value={infillPatchmatchDownscaleSize}
+      onChange={handleChange}
+      withInput
+      withSliderMarks
+      withReset
+      handleReset={handleReset}
+    />
+  );
+};
+
+export default memo(ParamInfillPatchmatchDownscaleSize);
diff --git a/invokeai/frontend/web/src/features/parameters/components/Parameters/Canvas/InfillAndScaling/ParamInfillTilesize.tsx b/invokeai/frontend/web/src/features/parameters/components/Parameters/Canvas/InfillAndScaling/ParamInfillTilesize.tsx
index fc6f02184c..4cb587767e 100644
--- a/invokeai/frontend/web/src/features/parameters/components/Parameters/Canvas/InfillAndScaling/ParamInfillTilesize.tsx
+++ b/invokeai/frontend/web/src/features/parameters/components/Parameters/Canvas/InfillAndScaling/ParamInfillTilesize.tsx
@@ -3,7 +3,7 @@ import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
 import IAISlider from 'common/components/IAISlider';
 import { generationSelector } from 'features/parameters/store/generationSelectors';
-import { setTileSize } from 'features/parameters/store/generationSlice';
+import { setInfillTileSize } from 'features/parameters/store/generationSlice';
 import { memo, useCallback } from 'react';
 
 import { useTranslation } from 'react-i18next';
@@ -11,10 +11,10 @@ import { useTranslation } from 'react-i18next';
 const selector = createSelector(
   [generationSelector],
   (parameters) => {
-    const { tileSize, infillMethod } = parameters;
+    const { infillTileSize, infillMethod } = parameters;
 
     return {
-      tileSize,
+      infillTileSize,
       infillMethod,
     };
   },
@@ -23,19 +23,19 @@ const selector = createSelector(
 
 const ParamInfillTileSize = () => {
   const dispatch = useAppDispatch();
-  const { tileSize, infillMethod } = useAppSelector(selector);
+  const { infillTileSize, infillMethod } = useAppSelector(selector);
 
   const { t } = useTranslation();
 
   const handleChange = useCallback(
     (v: number) => {
-      dispatch(setTileSize(v));
+      dispatch(setInfillTileSize(v));
     },
     [dispatch]
   );
 
   const handleReset = useCallback(() => {
-    dispatch(setTileSize(32));
+    dispatch(setInfillTileSize(32));
   }, [dispatch]);
 
   return (
@@ -45,7 +45,7 @@ const ParamInfillTileSize = () => {
       min={16}
       max={64}
       sliderNumberInputProps={{ max: 256 }}
-      value={tileSize}
+      value={infillTileSize}
       onChange={handleChange}
       withInput
       withSliderMarks
diff --git a/invokeai/frontend/web/src/features/parameters/store/generationSlice.ts b/invokeai/frontend/web/src/features/parameters/store/generationSlice.ts
index 5a830b6e95..7bbfc5149d 100644
--- a/invokeai/frontend/web/src/features/parameters/store/generationSlice.ts
+++ b/invokeai/frontend/web/src/features/parameters/store/generationSlice.ts
@@ -47,7 +47,8 @@ export interface GenerationState {
   shouldUseNoiseSettings: boolean;
   steps: StepsParam;
   threshold: number;
-  tileSize: number;
+  infillTileSize: number;
+  infillPatchmatchDownscaleSize: number;
   variationAmount: number;
   width: WidthParam;
   shouldUseSymmetry: boolean;
@@ -87,7 +88,8 @@ export const initialGenerationState: GenerationState = {
   shouldUseNoiseSettings: false,
   steps: 50,
   threshold: 0,
-  tileSize: 32,
+  infillTileSize: 32,
+  infillPatchmatchDownscaleSize: 2,
   variationAmount: 0.1,
   width: 512,
   shouldUseSymmetry: false,
@@ -212,12 +214,18 @@ export const generationSlice = createSlice({
     setCanvasCoherenceStrength: (state, action: PayloadAction<number>) => {
       state.canvasCoherenceStrength = action.payload;
     },
-    setTileSize: (state, action: PayloadAction<number>) => {
-      state.tileSize = action.payload;
-    },
     setInfillMethod: (state, action: PayloadAction<string>) => {
       state.infillMethod = action.payload;
     },
+    setInfillTileSize: (state, action: PayloadAction<number>) => {
+      state.infillTileSize = action.payload;
+    },
+    setInfillPatchmatchDownscaleSize: (
+      state,
+      action: PayloadAction<number>
+    ) => {
+      state.infillPatchmatchDownscaleSize = action.payload;
+    },
     setShouldUseSymmetry: (state, action: PayloadAction<boolean>) => {
       state.shouldUseSymmetry = action.payload;
     },
@@ -332,7 +340,8 @@ export const {
   setShouldRandomizeSeed,
   setSteps,
   setThreshold,
-  setTileSize,
+  setInfillTileSize,
+  setInfillPatchmatchDownscaleSize,
   setVariationAmount,
   setShouldUseSymmetry,
   setHorizontalSymmetrySteps,