Feat (ui): Add VAE Model to Recall Parameters (#5073)

* adding VAE recall when using all parameters * adding VAE to the RecallParameters tab in ImageMetadataActions * checking for nil vae and casting to null if undefined * adding default VAE to recall actions list if VAE is nullish * fix(ui): use `lodash-es` for tree-shakeable imports --------- Co-authored-by: psychedelicious <[email protected]>
invoke-ai · Nov 12, 2023 · 71e298b · 71e298b
1 parent 89a0394
commit 71e298b
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 1 deletion.
diff --git a/invokeai/frontend/web/public/locales/en.json b/invokeai/frontend/web/public/locales/en.json
@@ -583,6 +583,7 @@
         "strength": "Image to image strength",
         "Threshold": "Noise Threshold",
         "variations": "Seed-weight pairs",
+        "vae": "VAE", 
         "width": "Width",
         "workflow": "Workflow"
     },

diff --git a/...frontend/web/src/features/gallery/components/ImageMetadataViewer/ImageMetadataActions.tsx b/...frontend/web/src/features/gallery/components/ImageMetadataViewer/ImageMetadataActions.tsx
@@ -31,6 +31,7 @@ const ImageMetadataActions = (props: Props) => {
     recallCfgScale,
     recallModel,
     recallScheduler,
+    recallVaeModel,
     recallSteps,
     recallWidth,
     recallHeight,
@@ -72,6 +73,10 @@ const ImageMetadataActions = (props: Props) => {
     recallScheduler(metadata?.scheduler);
   }, [metadata?.scheduler, recallScheduler]);
 
+  const handleRecallVaeModel = useCallback(() => {
+    recallVaeModel(metadata?.vae);
+  }, [metadata?.vae, recallVaeModel]);
+
   const handleRecallSteps = useCallback(() => {
     recallSteps(metadata?.steps);
   }, [metadata?.steps, recallSteps]);
@@ -219,6 +224,11 @@ const ImageMetadataActions = (props: Props) => {
           onClick={handleRecallScheduler}
         />
       )}
+      <ImageMetadataItem
+        label={t('metadata.vae')}
+        value={metadata.vae?.model_name ?? 'Default'}
+        onClick={handleRecallVaeModel}
+      />
       {metadata.steps && (
         <ImageMetadataItem
           label={t('metadata.steps')}

diff --git a/invokeai/frontend/web/src/features/parameters/hooks/useRecallParameters.ts b/invokeai/frontend/web/src/features/parameters/hooks/useRecallParameters.ts
@@ -36,6 +36,7 @@ import {
   setRefinerStart,
   setRefinerSteps,
 } from 'features/sdxl/store/sdxlSlice';
+import { isNil } from 'lodash-es';
 import { useCallback } from 'react';
 import { useTranslation } from 'react-i18next';
 import { ImageDTO } from 'services/api/types';
@@ -65,8 +66,10 @@ import {
   setSeed,
   setSteps,
   setWidth,
+  vaeSelected,
 } from '../store/generationSlice';
 import {
+  isValidBoolean,
   isValidCfgScale,
   isValidControlNetModel,
   isValidHeight,
@@ -86,8 +89,8 @@ import {
   isValidSeed,
   isValidSteps,
   isValidStrength,
+  isValidVaeModel,
   isValidWidth,
-  isValidBoolean,
 } from '../types/parameterSchemas';
 
 const selector = createSelector(
@@ -306,6 +309,25 @@ export const useRecallParameters = () => {
     [dispatch, parameterSetToast, parameterNotSetToast]
   );
 
+  /**
+   * Recall vae model
+   */
+  const recallVaeModel = useCallback(
+    (vae: unknown) => {
+      if (!isValidVaeModel(vae) && !isNil(vae)) {
+        parameterNotSetToast();
+        return;
+      }
+      if (isNil(vae)) {
+        dispatch(vaeSelected(null));
+      } else {
+        dispatch(vaeSelected(vae));
+      }
+      parameterSetToast();
+    },
+    [dispatch, parameterSetToast, parameterNotSetToast]
+  );
+
   /**
    * Recall steps with toast
    */
@@ -757,6 +779,7 @@ export const useRecallParameters = () => {
         positive_prompt,
         negative_prompt,
         scheduler,
+        vae,
         seed,
         steps,
         width,
@@ -798,6 +821,13 @@ export const useRecallParameters = () => {
       if (isValidScheduler(scheduler)) {
         dispatch(setScheduler(scheduler));
       }
+      if (isValidVaeModel(vae) || isNil(vae)) {
+        if (isNil(vae)) {
+          dispatch(vaeSelected(null));
+        } else {
+          dispatch(vaeSelected(vae));
+        }
+      }
 
       if (isValidSeed(seed)) {
         dispatch(setSeed(seed));
@@ -932,6 +962,7 @@ export const useRecallParameters = () => {
     recallCfgScale,
     recallModel,
     recallScheduler,
+    recallVaeModel,
     recallSteps,
     recallWidth,
     recallHeight,