From 38e6e3b36bdbc9e1c595c5e7f3ba6845c3d56e94 Mon Sep 17 00:00:00 2001
From: blessedcoolant <54517381+blessedcoolant@users.noreply.github.com>
Date: Mon, 17 Jul 2023 16:07:38 +1200
Subject: [PATCH] feat: Add Quick Add To Scan Model

---
 .../AddModelsPanel/FoundModelsList.tsx        | 123 +++++++++++++++++-
 .../subpanels/AddModelsPanel/ScanModels.tsx   |   6 +-
 .../AddModelsPanel/SearchFolderForm.tsx       |   2 +-
 3 files changed, 121 insertions(+), 10 deletions(-)

diff --git a/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel/FoundModelsList.tsx b/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel/FoundModelsList.tsx
index c23d2bf9c8..58914fe04f 100644
--- a/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel/FoundModelsList.tsx
+++ b/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel/FoundModelsList.tsx
@@ -1,17 +1,72 @@
-import { Flex } from '@chakra-ui/react';
+import { Flex, Text } from '@chakra-ui/react';
+import { makeToast } from 'app/components/Toaster';
 import { RootState } from 'app/store/store';
-import { useAppSelector } from 'app/store/storeHooks';
-import { useGetModelsInFolderQuery } from 'services/api/endpoints/models';
+import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import IAIButton from 'common/components/IAIButton';
+import { addToast } from 'features/system/store/systemSlice';
+import { difference, map, values } from 'lodash-es';
+import { MouseEvent, useCallback } from 'react';
+import {
+  useGetMainModelsQuery,
+  useGetModelsInFolderQuery,
+  useImportMainModelsMutation,
+} from 'services/api/endpoints/models';
 
 export default function FoundModelsList() {
   const searchFolder = useAppSelector(
     (state: RootState) => state.modelmanager.searchFolder
   );
 
+  // Get all model paths from a given directory
   const { data: foundModels } = useGetModelsInFolderQuery({
     search_path: searchFolder ? searchFolder : '',
   });
 
+  // Get paths of models that are already installed
+  const { data: installedModels } = useGetMainModelsQuery();
+  const installedModelValues = values(installedModels?.entities);
+  const installedModelPaths = map(installedModelValues, 'path');
+
+  // Only select models those that aren't already installed to Invoke
+  const notInstalledModels = difference(foundModels, installedModelPaths);
+
+  const [importMainModel, { isLoading }] = useImportMainModelsMutation();
+  const dispatch = useAppDispatch();
+
+  const quickAddHandler = useCallback(
+    (e: MouseEvent<HTMLButtonElement>) => {
+      importMainModel({
+        body: {
+          location: e.currentTarget.id,
+        },
+      })
+        .unwrap()
+        .then((_) => {
+          dispatch(
+            addToast(
+              makeToast({
+                title: 'Added Model',
+                status: 'success',
+              })
+            )
+          );
+        })
+        .catch((error) => {
+          if (error) {
+            dispatch(
+              addToast(
+                makeToast({
+                  title: 'Faile To Add Model',
+                  status: 'error',
+                })
+              )
+            );
+          }
+        });
+    },
+    [dispatch, importMainModel]
+  );
+
   const renderFoundModels = () => {
     if (!searchFolder) return;
 
@@ -23,10 +78,68 @@ export default function FoundModelsList() {
       <Flex
         sx={{
           flexDirection: 'column',
+          gap: 2,
+          w: '100%',
         }}
       >
-        {foundModels.map((model) => (
-          <Flex key={model}>{model}</Flex>
+        <Flex p={2} gap={2}>
+          <Text
+            sx={{
+              fontWeight: 600,
+              color: 'accent.500',
+              _dark: {
+                color: 'accent.200',
+              },
+            }}
+          >
+            Found Models: {foundModels.length}
+          </Text>
+          <Text sx={{ fontWeight: 600 }}>
+            Not Installed: {notInstalledModels.length}
+          </Text>
+        </Flex>
+
+        {notInstalledModels.map((model) => (
+          <Flex
+            sx={{
+              p: 4,
+              gap: 4,
+              alignItems: 'center',
+              borderRadius: 4,
+              bg: 'base.200',
+              _dark: {
+                bg: 'base.800',
+              },
+            }}
+            key={model}
+          >
+            <Flex w="100%" sx={{ flexDirection: 'column' }}>
+              <Text sx={{ fontWeight: 600 }}>
+                {model.split('\\').slice(-1)[0]}
+              </Text>
+              <Text
+                sx={{
+                  fontSize: 'sm',
+                  color: 'base.600',
+                  _dark: {
+                    color: 'base.400',
+                  },
+                }}
+              >
+                {model}
+              </Text>
+            </Flex>
+            <Flex gap={2}>
+              <IAIButton
+                id={model}
+                onClick={quickAddHandler}
+                isLoading={isLoading}
+              >
+                Quick Add
+              </IAIButton>
+              <IAIButton>Advanced</IAIButton>
+            </Flex>
+          </Flex>
         ))}
       </Flex>
     );
diff --git a/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel/ScanModels.tsx b/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel/ScanModels.tsx
index d989463d75..e95264d2de 100644
--- a/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel/ScanModels.tsx
+++ b/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel/ScanModels.tsx
@@ -4,11 +4,9 @@ import SearchFolderForm from './SearchFolderForm';
 
 export default function ScanModels() {
   return (
-    <Flex flexDirection="column" w="100%">
+    <Flex flexDirection="column" w="100%" gap={2}>
       <SearchFolderForm />
-      <Flex
-        sx={{ maxHeight: window.innerHeight - 330, overflow: 'scroll', p: 4 }}
-      >
+      <Flex sx={{ maxHeight: window.innerHeight - 330, overflow: 'scroll' }}>
         <FoundModelsList />
       </Flex>
     </Flex>
diff --git a/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel/SearchFolderForm.tsx b/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel/SearchFolderForm.tsx
index 10d0f51665..6f17b62c66 100644
--- a/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel/SearchFolderForm.tsx
+++ b/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/AddModelsPanel/SearchFolderForm.tsx
@@ -47,7 +47,7 @@ function SearchFolderForm() {
           gap: 2,
           p: 4,
           borderRadius: 4,
-          background: 'base.400',
+          background: 'base.300',
           alignItems: 'center',
           _dark: {
             background: 'base.800',