diff --git a/invokeai/frontend/web/src/features/modelManagerV2/hooks/useInstallModel.ts b/invokeai/frontend/web/src/features/modelManagerV2/hooks/useInstallModel.ts new file mode 100644 index 0000000000..7636b9f314 --- /dev/null +++ b/invokeai/frontend/web/src/features/modelManagerV2/hooks/useInstallModel.ts @@ -0,0 +1,48 @@ +import { toast } from 'features/toast/toast'; +import { useCallback } from 'react'; +import { useTranslation } from 'react-i18next'; +import { useInstallModelMutation } from 'services/api/endpoints/models'; + +type InstallModelArg = { + source: string; + inplace?: boolean; + onSuccess?: () => void; + onError?: (error: unknown) => void; +}; + +export const useInstallModel = () => { + const { t } = useTranslation(); + const [_installModel, request] = useInstallModelMutation(); + + const installModel = useCallback( + ({ source, inplace, onSuccess, onError }: InstallModelArg) => { + _installModel({ source, inplace }) + .unwrap() + .then((_) => { + if (onSuccess) { + onSuccess(); + } + toast({ + id: 'MODEL_INSTALL_QUEUED', + title: t('toast.modelAddedSimple'), + status: 'success', + }); + }) + .catch((error) => { + if (onError) { + onError(error); + } + if (error) { + toast({ + id: 'MODEL_INSTALL_QUEUE_FAILED', + title: `${error.data.detail} `, + status: 'error', + }); + } + }); + }, + [_installModel, t] + ); + + return [installModel, request] as const; +}; diff --git a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/HuggingFaceFolder/HuggingFaceForm.tsx b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/HuggingFaceFolder/HuggingFaceForm.tsx index b3c456494f..ee5960f7d2 100644 --- a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/HuggingFaceFolder/HuggingFaceForm.tsx +++ b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/HuggingFaceFolder/HuggingFaceForm.tsx @@ -1,9 +1,9 @@ import { Button, Flex, FormControl, FormErrorMessage, FormHelperText, FormLabel, Input } from '@invoke-ai/ui-library'; -import { toast, ToastID } from 'features/toast/toast'; +import { useInstallModel } from 'features/modelManagerV2/hooks/useInstallModel'; import type { ChangeEventHandler } from 'react'; import { useCallback, useState } from 'react'; import { useTranslation } from 'react-i18next'; -import { useInstallModelMutation, useLazyGetHuggingFaceModelsQuery } from 'services/api/endpoints/models'; +import { useLazyGetHuggingFaceModelsQuery } from 'services/api/endpoints/models'; import { HuggingFaceResults } from './HuggingFaceResults'; @@ -14,41 +14,17 @@ export const HuggingFaceForm = () => { const { t } = useTranslation(); const [_getHuggingFaceModels, { isLoading, data }] = useLazyGetHuggingFaceModelsQuery(); - const [installModel] = useInstallModelMutation(); - - const handleInstallModel = useCallback( - (source: string) => { - installModel({ source }) - .unwrap() - .then((_) => { - toast({ - id: ToastID.MODEL_INSTALL_QUEUED, - title: t('toast.modelAddedSimple'), - status: 'success', - }); - }) - .catch((error) => { - if (error) { - toast({ - id: ToastID.MODEL_INSTALL_QUEUE_FAILED, - title: `${error.data.detail} `, - status: 'error', - }); - } - }); - }, - [installModel, t] - ); + const [installModel] = useInstallModel(); const getModels = useCallback(async () => { _getHuggingFaceModels(huggingFaceRepo) .unwrap() .then((response) => { if (response.is_diffusers) { - handleInstallModel(huggingFaceRepo); + installModel({ source: huggingFaceRepo }); setDisplayResults(false); } else if (response.urls?.length === 1 && response.urls[0]) { - handleInstallModel(response.urls[0]); + installModel({ source: response.urls[0] }); setDisplayResults(false); } else { setDisplayResults(true); @@ -57,7 +33,7 @@ export const HuggingFaceForm = () => { .catch((error) => { setErrorMessage(error.data.detail || ''); }); - }, [_getHuggingFaceModels, handleInstallModel, huggingFaceRepo]); + }, [_getHuggingFaceModels, installModel, huggingFaceRepo]); const handleSetHuggingFaceRepo: ChangeEventHandler = useCallback((e) => { setHuggingFaceRepo(e.target.value); diff --git a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/HuggingFaceFolder/HuggingFaceResultItem.tsx b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/HuggingFaceFolder/HuggingFaceResultItem.tsx index 225782aa3d..32970a3666 100644 --- a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/HuggingFaceFolder/HuggingFaceResultItem.tsx +++ b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/HuggingFaceFolder/HuggingFaceResultItem.tsx @@ -1,9 +1,8 @@ import { Flex, IconButton, Text } from '@invoke-ai/ui-library'; -import { toast, ToastID } from 'features/toast/toast'; +import { useInstallModel } from 'features/modelManagerV2/hooks/useInstallModel'; import { useCallback } from 'react'; import { useTranslation } from 'react-i18next'; import { PiPlusBold } from 'react-icons/pi'; -import { useInstallModelMutation } from 'services/api/endpoints/models'; type Props = { result: string; @@ -11,28 +10,11 @@ type Props = { export const HuggingFaceResultItem = ({ result }: Props) => { const { t } = useTranslation(); - const [installModel] = useInstallModelMutation(); + const [installModel] = useInstallModel(); - const handleInstall = useCallback(() => { - installModel({ source: result }) - .unwrap() - .then((_) => { - toast({ - id: ToastID.MODEL_INSTALL_QUEUED, - title: t('toast.modelAddedSimple'), - status: 'success', - }); - }) - .catch((error) => { - if (error) { - toast({ - id: ToastID.MODEL_INSTALL_QUEUE_FAILED, - title: `${error.data.detail} `, - status: 'error', - }); - } - }); - }, [installModel, result, t]); + const onClick = useCallback(() => { + installModel({ source: result }); + }, [installModel, result]); return ( @@ -42,7 +24,7 @@ export const HuggingFaceResultItem = ({ result }: Props) => { {result} - } onClick={handleInstall} size="sm" /> + } onClick={onClick} size="sm" /> ); }; diff --git a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/HuggingFaceFolder/HuggingFaceResults.tsx b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/HuggingFaceFolder/HuggingFaceResults.tsx index 3d78e76b83..826fd177ea 100644 --- a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/HuggingFaceFolder/HuggingFaceResults.tsx +++ b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/HuggingFaceFolder/HuggingFaceResults.tsx @@ -9,12 +9,11 @@ import { InputRightElement, } from '@invoke-ai/ui-library'; import ScrollableContent from 'common/components/OverlayScrollbars/ScrollableContent'; -import { toast, ToastID } from 'features/toast/toast'; +import { useInstallModel } from 'features/modelManagerV2/hooks/useInstallModel'; import type { ChangeEventHandler } from 'react'; import { useCallback, useMemo, useState } from 'react'; import { useTranslation } from 'react-i18next'; import { PiXBold } from 'react-icons/pi'; -import { useInstallModelMutation } from 'services/api/endpoints/models'; import { HuggingFaceResultItem } from './HuggingFaceResultItem'; @@ -26,7 +25,7 @@ export const HuggingFaceResults = ({ results }: HuggingFaceResultsProps) => { const { t } = useTranslation(); const [searchTerm, setSearchTerm] = useState(''); - const [installModel] = useInstallModelMutation(); + const [installModel] = useInstallModel(); const filteredResults = useMemo(() => { return results.filter((result) => { @@ -43,28 +42,11 @@ export const HuggingFaceResults = ({ results }: HuggingFaceResultsProps) => { setSearchTerm(''); }, []); - const handleAddAll = useCallback(() => { + const onClickAddAll = useCallback(() => { for (const result of filteredResults) { - installModel({ source: result }) - .unwrap() - .then((_) => { - toast({ - id: ToastID.MODEL_INSTALL_QUEUED, - title: t('toast.modelAddedSimple'), - status: 'success', - }); - }) - .catch((error) => { - if (error) { - toast({ - id: ToastID.MODEL_INSTALL_QUEUE_FAILED, - title: `${error.data.detail} `, - status: 'error', - }); - } - }); + installModel({ source: result }); } - }, [filteredResults, installModel, t]); + }, [filteredResults, installModel]); return ( <> @@ -73,7 +55,7 @@ export const HuggingFaceResults = ({ results }: HuggingFaceResultsProps) => { {t('modelManager.availableModels')} - diff --git a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/InstallModelForm.tsx b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/InstallModelForm.tsx index 8f1a6c38b8..cc052878bf 100644 --- a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/InstallModelForm.tsx +++ b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/InstallModelForm.tsx @@ -1,10 +1,9 @@ import { Button, Checkbox, Flex, FormControl, FormHelperText, FormLabel, Input } from '@invoke-ai/ui-library'; -import { toast, ToastID } from 'features/toast/toast'; +import { useInstallModel } from 'features/modelManagerV2/hooks/useInstallModel'; import { t } from 'i18next'; import { useCallback } from 'react'; import type { SubmitHandler } from 'react-hook-form'; import { useForm } from 'react-hook-form'; -import { useInstallModelMutation } from 'services/api/endpoints/models'; type SimpleImportModelConfig = { location: string; @@ -12,7 +11,7 @@ type SimpleImportModelConfig = { }; export const InstallModelForm = () => { - const [installModel, { isLoading }] = useInstallModelMutation(); + const [installModel, { isLoading }] = useInstallModel(); const { register, handleSubmit, formState, reset } = useForm({ defaultValues: { @@ -22,34 +21,22 @@ export const InstallModelForm = () => { mode: 'onChange', }); + const resetForm = useCallback(() => reset(undefined, { keepValues: true }), [reset]); + const onSubmit = useCallback>( (values) => { if (!values?.location) { return; } - installModel({ source: values.location, inplace: values.inplace }) - .unwrap() - .then((_) => { - toast({ - id: ToastID.MODEL_INSTALL_QUEUED, - title: t('toast.modelAddedSimple'), - status: 'success', - }); - reset(undefined, { keepValues: true }); - }) - .catch((error) => { - reset(undefined, { keepValues: true }); - if (error) { - toast({ - id: ToastID.MODEL_INSTALL_QUEUE_FAILED, - title: `${error.data.detail} `, - status: 'error', - }); - } - }); + installModel({ + source: values.location, + inplace: values.inplace, + onSuccess: resetForm, + onError: resetForm, + }); }, - [reset, installModel] + [installModel, resetForm] ); return ( diff --git a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/ScanFolder/ScanFolderResults.tsx b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/ScanFolder/ScanFolderResults.tsx index 67261730e5..749ef4c8e0 100644 --- a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/ScanFolder/ScanFolderResults.tsx +++ b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/ScanFolder/ScanFolderResults.tsx @@ -12,12 +12,12 @@ import { InputRightElement, } from '@invoke-ai/ui-library'; import ScrollableContent from 'common/components/OverlayScrollbars/ScrollableContent'; -import { toast, ToastID } from 'features/toast/toast'; +import { useInstallModel } from 'features/modelManagerV2/hooks/useInstallModel'; import type { ChangeEvent, ChangeEventHandler } from 'react'; import { useCallback, useMemo, useState } from 'react'; import { useTranslation } from 'react-i18next'; import { PiXBold } from 'react-icons/pi'; -import { type ScanFolderResponse, useInstallModelMutation } from 'services/api/endpoints/models'; +import type { ScanFolderResponse } from 'services/api/endpoints/models'; import { ScanModelResultItem } from './ScanFolderResultItem'; @@ -29,7 +29,7 @@ export const ScanModelsResults = ({ results }: ScanModelResultsProps) => { const { t } = useTranslation(); const [searchTerm, setSearchTerm] = useState(''); const [inplace, setInplace] = useState(true); - const [installModel] = useInstallModelMutation(); + const [installModel] = useInstallModel(); const filteredResults = useMemo(() => { return results.filter((result) => { @@ -55,49 +55,15 @@ export const ScanModelsResults = ({ results }: ScanModelResultsProps) => { if (result.is_installed) { continue; } - installModel({ source: result.path, inplace }) - .unwrap() - .then((_) => { - toast({ - id: ToastID.MODEL_INSTALL_QUEUED, - title: t('toast.modelAddedSimple'), - status: 'success', - }); - }) - .catch((error) => { - if (error) { - toast({ - id: ToastID.MODEL_INSTALL_QUEUE_FAILED, - title: `${error.data.detail} `, - status: 'error', - }); - } - }); + installModel({ source: result.path, inplace }); } - }, [filteredResults, installModel, inplace, t]); + }, [filteredResults, installModel, inplace]); const handleInstallOne = useCallback( (source: string) => { - installModel({ source, inplace }) - .unwrap() - .then((_) => { - toast({ - id: ToastID.MODEL_INSTALL_QUEUED, - title: t('toast.modelAddedSimple'), - status: 'success', - }); - }) - .catch((error) => { - if (error) { - toast({ - id: ToastID.MODEL_INSTALL_QUEUE_FAILED, - title: `${error.data.detail} `, - status: 'error', - }); - } - }); + installModel({ source, inplace }); }, - [installModel, inplace, t] + [installModel, inplace] ); return ( diff --git a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/StarterModels/StartModelsResultItem.tsx b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/StarterModels/StartModelsResultItem.tsx index 27a96c494c..98e1e39640 100644 --- a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/StarterModels/StartModelsResultItem.tsx +++ b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/StarterModels/StartModelsResultItem.tsx @@ -1,11 +1,10 @@ import { Badge, Box, Flex, IconButton, Text } from '@invoke-ai/ui-library'; +import { useInstallModel } from 'features/modelManagerV2/hooks/useInstallModel'; import ModelBaseBadge from 'features/modelManagerV2/subpanels/ModelManagerPanel/ModelBaseBadge'; -import { toast, ToastID } from 'features/toast/toast'; import { useCallback, useMemo } from 'react'; import { useTranslation } from 'react-i18next'; import { PiPlusBold } from 'react-icons/pi'; import type { GetStarterModelsResponse } from 'services/api/endpoints/models'; -import { useInstallModelMutation } from 'services/api/endpoints/models'; type Props = { result: GetStarterModelsResponse[number]; @@ -19,30 +18,13 @@ export const StarterModelsResultItem = ({ result }: Props) => { } return _allSources; }, [result]); - const [installModel] = useInstallModelMutation(); + const [installModel] = useInstallModel(); - const handleQuickAdd = useCallback(() => { + const onClick = useCallback(() => { for (const source of allSources) { - installModel({ source }) - .unwrap() - .then((_) => { - toast({ - id: ToastID.MODEL_INSTALL_QUEUED, - title: t('toast.modelAddedSimple'), - status: 'success', - }); - }) - .catch((error) => { - if (error) { - toast({ - id: ToastID.MODEL_INSTALL_QUEUE_FAILED, - title: `${error.data.detail} `, - status: 'error', - }); - } - }); + installModel({ source }); } - }, [allSources, installModel, t]); + }, [allSources, installModel]); return ( @@ -58,7 +40,7 @@ export const StarterModelsResultItem = ({ result }: Props) => { {result.is_installed ? ( {t('common.installed')} ) : ( - } onClick={handleQuickAdd} size="sm" /> + } onClick={onClick} size="sm" /> )}