From 92029e69c61b5a89b698eac3fd7d64de97b53566 Mon Sep 17 00:00:00 2001
From: blessedcoolant <54517381+blessedcoolant@users.noreply.github.com>
Date: Sun, 16 Jul 2023 15:48:44 +1200
Subject: [PATCH] feat: Update Checkpoint Model Edit to use config picker

---
 .../components/ImageGrid/GalleryImageGrid.tsx |  4 +-
 .../ModelManagerPanel/CheckpointModelEdit.tsx | 42 ++++++++++++++++---
 .../ModelManagerPanel/DiffusersModelEdit.tsx  | 11 ++++-
 3 files changed, 47 insertions(+), 10 deletions(-)

diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageGrid/GalleryImageGrid.tsx b/invokeai/frontend/web/src/features/gallery/components/ImageGrid/GalleryImageGrid.tsx
index 8b44b39ae9..7811e5f1c4 100644
--- a/invokeai/frontend/web/src/features/gallery/components/ImageGrid/GalleryImageGrid.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageGrid/GalleryImageGrid.tsx
@@ -16,14 +16,13 @@ import {
   ASSETS_CATEGORIES,
   IMAGE_CATEGORIES,
   IMAGE_LIMIT,
-  selectImagesAll,
 } from 'features/gallery//store/gallerySlice';
 import { selectFilteredImages } from 'features/gallery/store/gallerySelectors';
 import { VirtuosoGrid } from 'react-virtuoso';
 import { receivedPageOfImages } from 'services/api/thunks/image';
+import { useListBoardImagesQuery } from '../../../../services/api/endpoints/boardImages';
 import ImageGridItemContainer from './ImageGridItemContainer';
 import ImageGridListContainer from './ImageGridListContainer';
-import { useListBoardImagesQuery } from '../../../../services/api/endpoints/boardImages';
 
 const selector = createSelector(
   [stateSelector, selectFilteredImages],
@@ -180,7 +179,6 @@ const GalleryImageGrid = () => {
       </Box>
     );
   }
-  console.log({ selectedBoardId });
 
   if (status !== 'rejected') {
     return (
diff --git a/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/ModelManagerPanel/CheckpointModelEdit.tsx b/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/ModelManagerPanel/CheckpointModelEdit.tsx
index 6078ba4a68..960fa1cbf5 100644
--- a/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/ModelManagerPanel/CheckpointModelEdit.tsx
+++ b/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/ModelManagerPanel/CheckpointModelEdit.tsx
@@ -4,17 +4,20 @@ import { makeToast } from 'app/components/Toaster';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import IAIButton from 'common/components/IAIButton';
 import IAIMantineTextInput from 'common/components/IAIMantineInput';
+import IAISimpleCheckbox from 'common/components/IAISimpleCheckbox';
 import { MODEL_TYPE_MAP } from 'features/parameters/types/constants';
 import { selectIsBusy } from 'features/system/store/systemSelectors';
 import { addToast } from 'features/system/store/systemSlice';
-import { useCallback } from 'react';
+import { useCallback, useEffect, useState } from 'react';
 import { useTranslation } from 'react-i18next';
 import {
   CheckpointModelConfigEntity,
+  useGetCheckpointConfigsQuery,
   useUpdateMainModelsMutation,
 } from 'services/api/endpoints/models';
 import { CheckpointModelConfig } from 'services/api/types';
 import BaseModelSelect from '../shared/BaseModelSelect';
+import CheckpointConfigsSelect from '../shared/CheckpointConfigsSelect';
 import ModelVariantSelect from '../shared/ModelVariantSelect';
 import ModelConvert from './ModelConvert';
 
@@ -28,6 +31,15 @@ export default function CheckpointModelEdit(props: CheckpointModelEditProps) {
   const { model } = props;
 
   const [updateMainModel, { isLoading }] = useUpdateMainModelsMutation();
+  const { data: availableCheckpointConfigs } = useGetCheckpointConfigsQuery();
+
+  const [useCustomConfig, setUseCustomConfig] = useState<boolean>(false);
+
+  useEffect(() => {
+    if (!availableCheckpointConfigs?.includes(model.config)) {
+      setUseCustomConfig(true);
+    }
+  }, [availableCheckpointConfigs, model.config]);
 
   const dispatch = useAppDispatch();
   const { t } = useTranslation();
@@ -123,12 +135,15 @@ export default function CheckpointModelEdit(props: CheckpointModelEditProps) {
               {...checkpointEditForm.getInputProps('description')}
             />
             <BaseModelSelect
+              required
               {...checkpointEditForm.getInputProps('base_model')}
             />
             <ModelVariantSelect
+              required
               {...checkpointEditForm.getInputProps('variant')}
             />
             <IAIMantineTextInput
+              required
               label={t('modelManager.modelLocation')}
               {...checkpointEditForm.getInputProps('path')}
             />
@@ -136,10 +151,27 @@ export default function CheckpointModelEdit(props: CheckpointModelEditProps) {
               label={t('modelManager.vaeLocation')}
               {...checkpointEditForm.getInputProps('vae')}
             />
-            <IAIMantineTextInput
-              label={t('modelManager.config')}
-              {...checkpointEditForm.getInputProps('config')}
-            />
+
+            <Flex flexDirection="column" gap={2}>
+              {!useCustomConfig ? (
+                <CheckpointConfigsSelect
+                  required
+                  {...checkpointEditForm.getInputProps('config')}
+                />
+              ) : (
+                <IAIMantineTextInput
+                  required
+                  label={t('modelManager.config')}
+                  {...checkpointEditForm.getInputProps('config')}
+                />
+              )}
+              <IAISimpleCheckbox
+                isChecked={useCustomConfig}
+                onChange={() => setUseCustomConfig(!useCustomConfig)}
+                label="Use Custom Config"
+              />
+            </Flex>
+
             <IAIButton
               type="submit"
               isDisabled={isBusy || isLoading}
diff --git a/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/ModelManagerPanel/DiffusersModelEdit.tsx b/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/ModelManagerPanel/DiffusersModelEdit.tsx
index ed8f64caed..9da47a7397 100644
--- a/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/ModelManagerPanel/DiffusersModelEdit.tsx
+++ b/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/ModelManagerPanel/DiffusersModelEdit.tsx
@@ -112,9 +112,16 @@ export default function DiffusersModelEdit(props: DiffusersModelEditProps) {
             label={t('modelManager.description')}
             {...diffusersEditForm.getInputProps('description')}
           />
-          <BaseModelSelect {...diffusersEditForm.getInputProps('base_model')} />
-          <ModelVariantSelect {...diffusersEditForm.getInputProps('variant')} />
+          <BaseModelSelect
+            required
+            {...diffusersEditForm.getInputProps('base_model')}
+          />
+          <ModelVariantSelect
+            required
+            {...diffusersEditForm.getInputProps('variant')}
+          />
           <IAIMantineTextInput
+            required
             label={t('modelManager.modelLocation')}
             {...diffusersEditForm.getInputProps('path')}
           />