From 6fa244a343d9c6b6f8737f85543f737e2c4d5dc3 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Wed, 26 Jul 2023 00:39:58 +1000
Subject: [PATCH] feat(ui): add vae precision select

---
 .../buildLinearImageToImageGraph.ts           |  3 ++
 .../buildLinearTextToImageGraph.ts            |  2 +
 .../Core/ParamModelandVAEandScheduler.tsx     | 20 ++++----
 .../Parameters/VAEModel/ParamVAEPrecision.tsx | 46 +++++++++++++++++++
 4 files changed, 61 insertions(+), 10 deletions(-)
 create mode 100644 invokeai/frontend/web/src/features/parameters/components/Parameters/VAEModel/ParamVAEPrecision.tsx

diff --git a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearImageToImageGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearImageToImageGraph.ts
index f5f5e670c9..856ed89e5a 100644
--- a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearImageToImageGraph.ts
+++ b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearImageToImageGraph.ts
@@ -46,6 +46,7 @@ export const buildLinearImageToImageGraph = (
     clipSkip,
     shouldUseCpuNoise,
     shouldUseNoiseSettings,
+    vaePrecision,
   } = state.generation;
 
   // TODO: add batch functionality
@@ -113,6 +114,7 @@ export const buildLinearImageToImageGraph = (
       [LATENTS_TO_IMAGE]: {
         type: 'l2i',
         id: LATENTS_TO_IMAGE,
+        fp32: vaePrecision === 'fp32' ? true : false,
       },
       [LATENTS_TO_LATENTS]: {
         type: 'l2l',
@@ -129,6 +131,7 @@ export const buildLinearImageToImageGraph = (
         // image: {
         //   image_name: initialImage.image_name,
         // },
+        fp32: vaePrecision === 'fp32' ? true : false,
       },
     },
     edges: [
diff --git a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearTextToImageGraph.ts b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearTextToImageGraph.ts
index c8529f95d3..7b36f6b391 100644
--- a/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearTextToImageGraph.ts
+++ b/invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearTextToImageGraph.ts
@@ -34,6 +34,7 @@ export const buildLinearTextToImageGraph = (
     clipSkip,
     shouldUseCpuNoise,
     shouldUseNoiseSettings,
+    vaePrecision,
   } = state.generation;
 
   const use_cpu = shouldUseNoiseSettings
@@ -95,6 +96,7 @@ export const buildLinearTextToImageGraph = (
       [LATENTS_TO_IMAGE]: {
         type: 'l2i',
         id: LATENTS_TO_IMAGE,
+        fp32: vaePrecision === 'fp32' ? true : false,
       },
     },
     edges: [
diff --git a/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamModelandVAEandScheduler.tsx b/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamModelandVAEandScheduler.tsx
index 74418de1d3..37c5eb30c4 100644
--- a/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamModelandVAEandScheduler.tsx
+++ b/invokeai/frontend/web/src/features/parameters/components/Parameters/Core/ParamModelandVAEandScheduler.tsx
@@ -4,6 +4,7 @@ import { memo } from 'react';
 import ParamMainModelSelect from '../MainModel/ParamMainModelSelect';
 import ParamVAEModelSelect from '../VAEModel/ParamVAEModelSelect';
 import ParamScheduler from './ParamScheduler';
+import ParamVAEPrecision from '../VAEModel/ParamVAEPrecision';
 
 const ParamModelandVAEandScheduler = () => {
   const isVaeEnabled = useFeatureStatus('vae').isFeatureEnabled;
@@ -13,16 +14,15 @@ const ParamModelandVAEandScheduler = () => {
       <Box w="full">
         <ParamMainModelSelect />
       </Box>
-      <Flex gap={3} w="full">
-        {isVaeEnabled && (
-          <Box w="full">
-            <ParamVAEModelSelect />
-          </Box>
-        )}
-        <Box w="full">
-          <ParamScheduler />
-        </Box>
-      </Flex>
+      <Box w="full">
+        <ParamScheduler />
+      </Box>
+      {isVaeEnabled && (
+        <Flex w="full" gap={3}>
+          <ParamVAEModelSelect />
+          <ParamVAEPrecision />
+        </Flex>
+      )}
     </Flex>
   );
 };
diff --git a/invokeai/frontend/web/src/features/parameters/components/Parameters/VAEModel/ParamVAEPrecision.tsx b/invokeai/frontend/web/src/features/parameters/components/Parameters/VAEModel/ParamVAEPrecision.tsx
new file mode 100644
index 0000000000..c57cdc1132
--- /dev/null
+++ b/invokeai/frontend/web/src/features/parameters/components/Parameters/VAEModel/ParamVAEPrecision.tsx
@@ -0,0 +1,46 @@
+import { createSelector } from '@reduxjs/toolkit';
+import { stateSelector } from 'app/store/store';
+import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
+import IAIMantineSelect from 'common/components/IAIMantineSelect';
+import { vaePrecisionChanged } from 'features/parameters/store/generationSlice';
+import { PrecisionParam } from 'features/parameters/types/parameterSchemas';
+import { memo, useCallback } from 'react';
+
+const selector = createSelector(
+  stateSelector,
+  ({ generation }) => {
+    const { vaePrecision } = generation;
+    return { vaePrecision };
+  },
+  defaultSelectorOptions
+);
+
+const DATA = ['fp16', 'fp32'];
+
+const ParamVAEModelSelect = () => {
+  const dispatch = useAppDispatch();
+  const { vaePrecision } = useAppSelector(selector);
+
+  const handleChange = useCallback(
+    (v: string | null) => {
+      if (!v) {
+        return;
+      }
+
+      dispatch(vaePrecisionChanged(v as PrecisionParam));
+    },
+    [dispatch]
+  );
+
+  return (
+    <IAIMantineSelect
+      label="VAE Precision"
+      value={vaePrecision}
+      data={DATA}
+      onChange={handleChange}
+    />
+  );
+};
+
+export default memo(ParamVAEModelSelect);