From cd033f4ead1a9eb4be35cff6118c16860d1accfa Mon Sep 17 00:00:00 2001
From: blessedcoolant <54517381+blessedcoolant@users.noreply.github.com>
Date: Sun, 16 Jul 2023 01:57:42 +1200
Subject: [PATCH] fix: Refine some UI

---
 .../ModelManager/subpanels/AddModelsPanel.tsx |   2 +
 .../subpanels/AddModelsPanel/AddModels.tsx    | 136 +++++-------------
 .../AddModelsPanel/AutoAddModels.tsx          | 108 ++++++++++++++
 3 files changed, 146 insertions(+), 100 deletions(-)
 create mode 100644 invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel/AutoAddModels.tsx

diff --git a/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel.tsx b/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel.tsx
index d868dc6c80..cac6c7a136 100644
--- a/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel.tsx
+++ b/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel.tsx
@@ -18,6 +18,7 @@ export default function AddModelsPanel() {
           onClick={() => setAddModelTab('add')}
           isChecked={addModelTab == 'add'}
           size="sm"
+          width="100%"
         >
           {t('modelManager.addModel')}
         </IAIButton>
@@ -25,6 +26,7 @@ export default function AddModelsPanel() {
           onClick={() => setAddModelTab('scan')}
           isChecked={addModelTab == 'scan'}
           size="sm"
+          width="100%"
         >
           {t('modelManager.scanForModels')}
         </IAIButton>
diff --git a/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel/AddModels.tsx b/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel/AddModels.tsx
index e4d6ffb663..e667c17d3e 100644
--- a/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel/AddModels.tsx
+++ b/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel/AddModels.tsx
@@ -1,111 +1,47 @@
-import { Flex } from '@chakra-ui/react';
-// import { addNewModel } from 'app/socketio/actions';
-import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
-import { useTranslation } from 'react-i18next';
-
-import { SelectItem } from '@mantine/core';
-import { useForm } from '@mantine/form';
-import { makeToast } from 'app/components/Toaster';
-import { RootState } from 'app/store/store';
+import { ButtonGroup, Flex } from '@chakra-ui/react';
 import IAIButton from 'common/components/IAIButton';
-import IAIMantineTextInput from 'common/components/IAIMantineInput';
-import IAIMantineSelect from 'common/components/IAIMantineSelect';
-import { addToast } from 'features/system/store/systemSlice';
-import { useImportMainModelsMutation } from 'services/api/endpoints/models';
-
-const predictionSelectData: SelectItem[] = [
-  { label: 'None', value: 'none' },
-  { label: 'v_prediction', value: 'v_prediction' },
-  { label: 'epsilon', value: 'epsilon' },
-  { label: 'sample', value: 'sample' },
-];
-
-type ExtendedImportModelConfig = {
-  location: string;
-  prediction_type?: 'v_prediction' | 'epsilon' | 'sample' | 'none' | undefined;
-};
+import { useState } from 'react';
+import AutoAddModels from './AutoAddModels';
 
 export default function AddModels() {
-  const dispatch = useAppDispatch();
-  const { t } = useTranslation();
-
-  const isProcessing = useAppSelector(
-    (state: RootState) => state.system.isProcessing
+  const [addModelMode, setAddModelMode] = useState<'simple' | 'advanced'>(
+    'simple'
   );
-
-  const [importMainModel, { isLoading }] = useImportMainModelsMutation();
-
-  const addModelForm = useForm<ExtendedImportModelConfig>({
-    initialValues: {
-      location: '',
-      prediction_type: undefined,
-    },
-  });
-
-  const handleAddModelSubmit = (values: ExtendedImportModelConfig) => {
-    const importModelResponseBody = {
-      location: values.location,
-      prediction_type:
-        values.prediction_type === 'none' ? undefined : values.prediction_type,
-    };
-
-    importMainModel({ body: importModelResponseBody })
-      .unwrap()
-      .then((_) => {
-        dispatch(
-          addToast(
-            makeToast({
-              title: 'Model Added',
-              status: 'success',
-            })
-          )
-        );
-        addModelForm.reset();
-      })
-      .catch((error) => {
-        if (error) {
-          console.log(error);
-          dispatch(
-            addToast(
-              makeToast({
-                title: `${error.data.detail} `,
-                status: 'error',
-              })
-            )
-          );
-        }
-      });
-  };
-
   return (
-    <form onSubmit={addModelForm.onSubmit((v) => handleAddModelSubmit(v))}>
-      <Flex
-        flexDirection="column"
-        overflow="scroll"
-        maxHeight={window.innerHeight - 270}
-        width="100%"
-        gap={4}
-      >
-        <IAIMantineTextInput
-          label="Model Location"
-          placeholder="Provide a path to a local Diffusers model, local checkpoint / safetensors model or a HuggingFace Repo ID"
-          w="100%"
-          {...addModelForm.getInputProps('location')}
-        />
-        <IAIMantineSelect
-          label="Prediction Type (for Stable Diffusion 2.x Models only)"
-          data={predictionSelectData}
-          defaultValue="none"
-          {...addModelForm.getInputProps('prediction_type')}
-        />
+    <Flex
+      flexDirection="column"
+      width="100%"
+      overflow="scroll"
+      maxHeight={window.innerHeight - 270}
+      gap={4}
+    >
+      <ButtonGroup isAttached>
         <IAIButton
-          type="submit"
-          isLoading={isLoading}
-          isDisabled={isLoading || isProcessing}
+          size="sm"
+          isChecked={addModelMode == 'simple'}
+          onClick={() => setAddModelMode('simple')}
         >
-          {t('modelManager.addModel')}
+          Simple
         </IAIButton>
+        <IAIButton
+          size="sm"
+          isChecked={addModelMode == 'advanced'}
+          onClick={() => setAddModelMode('advanced')}
+        >
+          Advanced
+        </IAIButton>
+      </ButtonGroup>
+      <Flex
+        sx={{
+          p: 4,
+          borderRadius: 4,
+          background: 'base.200',
+          _dark: { background: 'base.800' },
+        }}
+      >
+        {addModelMode === 'simple' && <AutoAddModels />}
+        {addModelMode === 'advanced' && null}
       </Flex>
-    </form>
+    </Flex>
   );
 }
diff --git a/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel/AutoAddModels.tsx b/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel/AutoAddModels.tsx
new file mode 100644
index 0000000000..55a33ca1b8
--- /dev/null
+++ b/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel/AutoAddModels.tsx
@@ -0,0 +1,108 @@
+import { Flex } from '@chakra-ui/react';
+// import { addNewModel } from 'app/socketio/actions';
+import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { useTranslation } from 'react-i18next';
+
+import { SelectItem } from '@mantine/core';
+import { useForm } from '@mantine/form';
+import { makeToast } from 'app/components/Toaster';
+import { RootState } from 'app/store/store';
+import IAIButton from 'common/components/IAIButton';
+import IAIMantineTextInput from 'common/components/IAIMantineInput';
+import IAIMantineSelect from 'common/components/IAIMantineSelect';
+import { addToast } from 'features/system/store/systemSlice';
+import { useImportMainModelsMutation } from 'services/api/endpoints/models';
+
+const predictionSelectData: SelectItem[] = [
+  { label: 'None', value: 'none' },
+  { label: 'v_prediction', value: 'v_prediction' },
+  { label: 'epsilon', value: 'epsilon' },
+  { label: 'sample', value: 'sample' },
+];
+
+type ExtendedImportModelConfig = {
+  location: string;
+  prediction_type?: 'v_prediction' | 'epsilon' | 'sample' | 'none' | undefined;
+};
+
+export default function AutoAddModels() {
+  const dispatch = useAppDispatch();
+  const { t } = useTranslation();
+
+  const isProcessing = useAppSelector(
+    (state: RootState) => state.system.isProcessing
+  );
+
+  const [importMainModel, { isLoading }] = useImportMainModelsMutation();
+
+  const addModelForm = useForm<ExtendedImportModelConfig>({
+    initialValues: {
+      location: '',
+      prediction_type: undefined,
+    },
+  });
+
+  const handleAddModelSubmit = (values: ExtendedImportModelConfig) => {
+    const importModelResponseBody = {
+      location: values.location,
+      prediction_type:
+        values.prediction_type === 'none' ? undefined : values.prediction_type,
+    };
+
+    importMainModel({ body: importModelResponseBody })
+      .unwrap()
+      .then((_) => {
+        dispatch(
+          addToast(
+            makeToast({
+              title: 'Model Added',
+              status: 'success',
+            })
+          )
+        );
+        addModelForm.reset();
+      })
+      .catch((error) => {
+        if (error) {
+          console.log(error);
+          dispatch(
+            addToast(
+              makeToast({
+                title: `${error.data.detail} `,
+                status: 'error',
+              })
+            )
+          );
+        }
+      });
+  };
+
+  return (
+    <form
+      onSubmit={addModelForm.onSubmit((v) => handleAddModelSubmit(v))}
+      style={{ width: '100%' }}
+    >
+      <Flex flexDirection="column" width="100%" gap={4}>
+        <IAIMantineTextInput
+          label="Model Location"
+          placeholder="Provide a path to a local Diffusers model, local checkpoint / safetensors model or a HuggingFace Repo ID"
+          w="100%"
+          {...addModelForm.getInputProps('location')}
+        />
+        <IAIMantineSelect
+          label="Prediction Type (for Stable Diffusion 2.x Models only)"
+          data={predictionSelectData}
+          defaultValue="none"
+          {...addModelForm.getInputProps('prediction_type')}
+        />
+        <IAIButton
+          type="submit"
+          isLoading={isLoading}
+          isDisabled={isLoading || isProcessing}
+        >
+          {t('modelManager.addModel')}
+        </IAIButton>
+      </Flex>
+    </form>
+  );
+}