Skip to content

Commit f2b9684

Browse files
psychedeliciouspsychedelicious
committed
tidy(ui): split install model into helper hook
This was duplicated like 7 times or so
1 parent a66b349 commit f2b9684

File tree

7 files changed

+90
-167
lines changed

7 files changed

+90
-167
lines changed

‎invokeai/frontend/web/src/features/modelManagerV2/hooks/useInstallModel.ts

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
import { toast } from 'features/toast/toast';
2+
import { useCallback } from 'react';
3+
import { useTranslation } from 'react-i18next';
4+
import { useInstallModelMutation } from 'services/api/endpoints/models';
5+
6+
type InstallModelArg = {
7+
source: string;
8+
inplace?: boolean;
9+
onSuccess?: () => void;
10+
onError?: (error: unknown) => void;
11+
};
12+
13+
export const useInstallModel = () => {
14+
const { t } = useTranslation();
15+
const [_installModel, request] = useInstallModelMutation();
16+
17+
const installModel = useCallback(
18+
({ source, inplace, onSuccess, onError }: InstallModelArg) => {
19+
_installModel({ source, inplace })
20+
.unwrap()
21+
.then((_) => {
22+
if (onSuccess) {
23+
onSuccess();
24+
}
25+
toast({
26+
id: 'MODEL_INSTALL_QUEUED',
27+
title: t('toast.modelAddedSimple'),
28+
status: 'success',
29+
});
30+
})
31+
.catch((error) => {
32+
if (onError) {
33+
onError(error);
34+
}
35+
if (error) {
36+
toast({
37+
id: 'MODEL_INSTALL_QUEUE_FAILED',
38+
title: `${error.data.detail} `,
39+
status: 'error',
40+
});
41+
}
42+
});
43+
},
44+
[_installModel, t]
45+
);
46+
47+
return [installModel, request] as const;
48+
};

‎invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/HuggingFaceFolder/HuggingFaceForm.tsx

Lines changed: 6 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
import { Button, Flex, FormControl, FormErrorMessage, FormHelperText, FormLabel, Input } from '@invoke-ai/ui-library';
2-
import { toast, ToastID } from 'features/toast/toast';
2+
import { useInstallModel } from 'features/modelManagerV2/hooks/useInstallModel';
33
import type { ChangeEventHandler } from 'react';
44
import { useCallback, useState } from 'react';
55
import { useTranslation } from 'react-i18next';
6-
import { useInstallModelMutation, useLazyGetHuggingFaceModelsQuery } from 'services/api/endpoints/models';
6+
import { useLazyGetHuggingFaceModelsQuery } from 'services/api/endpoints/models';
77

88
import { HuggingFaceResults } from './HuggingFaceResults';
99

@@ -14,41 +14,17 @@ export const HuggingFaceForm = () => {
1414
const { t } = useTranslation();
1515

1616
const [_getHuggingFaceModels, { isLoading, data }] = useLazyGetHuggingFaceModelsQuery();
17-
const [installModel] = useInstallModelMutation();
18-
19-
const handleInstallModel = useCallback(
20-
(source: string) => {
21-
installModel({ source })
22-
.unwrap()
23-
.then((_) => {
24-
toast({
25-
id: ToastID.MODEL_INSTALL_QUEUED,
26-
title: t('toast.modelAddedSimple'),
27-
status: 'success',
28-
});
29-
})
30-
.catch((error) => {
31-
if (error) {
32-
toast({
33-
id: ToastID.MODEL_INSTALL_QUEUE_FAILED,
34-
title: `${error.data.detail} `,
35-
status: 'error',
36-
});
37-
}
38-
});
39-
},
40-
[installModel, t]
41-
);
17+
const [installModel] = useInstallModel();
4218

4319
const getModels = useCallback(async () => {
4420
_getHuggingFaceModels(huggingFaceRepo)
4521
.unwrap()
4622
.then((response) => {
4723
if (response.is_diffusers) {
48-
handleInstallModel(huggingFaceRepo);
24+
installModel({ source: huggingFaceRepo });
4925
setDisplayResults(false);
5026
} else if (response.urls?.length === 1 && response.urls[0]) {
51-
handleInstallModel(response.urls[0]);
27+
installModel({ source: response.urls[0] });
5228
setDisplayResults(false);
5329
} else {
5430
setDisplayResults(true);
@@ -57,7 +33,7 @@ export const HuggingFaceForm = () => {
5733
.catch((error) => {
5834
setErrorMessage(error.data.detail || '');
5935
});
60-
}, [_getHuggingFaceModels, handleInstallModel, huggingFaceRepo]);
36+
}, [_getHuggingFaceModels, installModel, huggingFaceRepo]);
6137

6238
const handleSetHuggingFaceRepo: ChangeEventHandler<HTMLInputElement> = useCallback((e) => {
6339
setHuggingFaceRepo(e.target.value);

‎invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/HuggingFaceFolder/HuggingFaceResultItem.tsx

Lines changed: 6 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,38 +1,20 @@
11
import { Flex, IconButton, Text } from '@invoke-ai/ui-library';
2-
import { toast, ToastID } from 'features/toast/toast';
2+
import { useInstallModel } from 'features/modelManagerV2/hooks/useInstallModel';
33
import { useCallback } from 'react';
44
import { useTranslation } from 'react-i18next';
55
import { PiPlusBold } from 'react-icons/pi';
6-
import { useInstallModelMutation } from 'services/api/endpoints/models';
76

87
type Props = {
98
result: string;
109
};
1110
export const HuggingFaceResultItem = ({ result }: Props) => {
1211
const { t } = useTranslation();
1312

14-
const [installModel] = useInstallModelMutation();
13+
const [installModel] = useInstallModel();
1514

16-
const handleInstall = useCallback(() => {
17-
installModel({ source: result })
18-
.unwrap()
19-
.then((_) => {
20-
toast({
21-
id: ToastID.MODEL_INSTALL_QUEUED,
22-
title: t('toast.modelAddedSimple'),
23-
status: 'success',
24-
});
25-
})
26-
.catch((error) => {
27-
if (error) {
28-
toast({
29-
id: ToastID.MODEL_INSTALL_QUEUE_FAILED,
30-
title: `${error.data.detail} `,
31-
status: 'error',
32-
});
33-
}
34-
});
35-
}, [installModel, result, t]);
15+
const onClick = useCallback(() => {
16+
installModel({ source: result });
17+
}, [installModel, result]);
3618

3719
return (
3820
<Flex alignItems="center" justifyContent="space-between" w="100%" gap={3}>
@@ -42,7 +24,7 @@ export const HuggingFaceResultItem = ({ result }: Props) => {
4224
{result}
4325
</Text>
4426
</Flex>
45-
<IconButton aria-label={t('modelManager.install')} icon={<PiPlusBold />} onClick={handleInstall} size="sm" />
27+
<IconButton aria-label={t('modelManager.install')} icon={<PiPlusBold />} onClick={onClick} size="sm" />
4628
</Flex>
4729
);
4830
};

‎invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/HuggingFaceFolder/HuggingFaceResults.tsx

Lines changed: 6 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,11 @@ import {
99
InputRightElement,
1010
} from '@invoke-ai/ui-library';
1111
import ScrollableContent from 'common/components/OverlayScrollbars/ScrollableContent';
12-
import { toast, ToastID } from 'features/toast/toast';
12+
import { useInstallModel } from 'features/modelManagerV2/hooks/useInstallModel';
1313
import type { ChangeEventHandler } from 'react';
1414
import { useCallback, useMemo, useState } from 'react';
1515
import { useTranslation } from 'react-i18next';
1616
import { PiXBold } from 'react-icons/pi';
17-
import { useInstallModelMutation } from 'services/api/endpoints/models';
1817

1918
import { HuggingFaceResultItem } from './HuggingFaceResultItem';
2019

@@ -26,7 +25,7 @@ export const HuggingFaceResults = ({ results }: HuggingFaceResultsProps) => {
2625
const { t } = useTranslation();
2726
const [searchTerm, setSearchTerm] = useState('');
2827

29-
const [installModel] = useInstallModelMutation();
28+
const [installModel] = useInstallModel();
3029

3130
const filteredResults = useMemo(() => {
3231
return results.filter((result) => {
@@ -43,28 +42,11 @@ export const HuggingFaceResults = ({ results }: HuggingFaceResultsProps) => {
4342
setSearchTerm('');
4443
}, []);
4544

46-
const handleAddAll = useCallback(() => {
45+
const onClickAddAll = useCallback(() => {
4746
for (const result of filteredResults) {
48-
installModel({ source: result })
49-
.unwrap()
50-
.then((_) => {
51-
toast({
52-
id: ToastID.MODEL_INSTALL_QUEUED,
53-
title: t('toast.modelAddedSimple'),
54-
status: 'success',
55-
});
56-
})
57-
.catch((error) => {
58-
if (error) {
59-
toast({
60-
id: ToastID.MODEL_INSTALL_QUEUE_FAILED,
61-
title: `${error.data.detail} `,
62-
status: 'error',
63-
});
64-
}
65-
});
47+
installModel({ source: result });
6648
}
67-
}, [filteredResults, installModel, t]);
49+
}, [filteredResults, installModel]);
6850

6951
return (
7052
<>
@@ -73,7 +55,7 @@ export const HuggingFaceResults = ({ results }: HuggingFaceResultsProps) => {
7355
<Flex justifyContent="space-between" alignItems="center">
7456
<Heading size="sm">{t('modelManager.availableModels')}</Heading>
7557
<Flex alignItems="center" gap={3}>
76-
<Button size="sm" onClick={handleAddAll} isDisabled={results.length === 0} flexShrink={0}>
58+
<Button size="sm" onClick={onClickAddAll} isDisabled={results.length === 0} flexShrink={0}>
7759
{t('modelManager.installAll')}
7860
</Button>
7961
<InputGroup w={64} size="xs">

‎invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/InstallModelForm.tsx

Lines changed: 11 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,17 @@
11
import { Button, Checkbox, Flex, FormControl, FormHelperText, FormLabel, Input } from '@invoke-ai/ui-library';
2-
import { toast, ToastID } from 'features/toast/toast';
2+
import { useInstallModel } from 'features/modelManagerV2/hooks/useInstallModel';
33
import { t } from 'i18next';
44
import { useCallback } from 'react';
55
import type { SubmitHandler } from 'react-hook-form';
66
import { useForm } from 'react-hook-form';
7-
import { useInstallModelMutation } from 'services/api/endpoints/models';
87

98
type SimpleImportModelConfig = {
109
location: string;
1110
inplace: boolean;
1211
};
1312

1413
export const InstallModelForm = () => {
15-
const [installModel, { isLoading }] = useInstallModelMutation();
14+
const [installModel, { isLoading }] = useInstallModel();
1615

1716
const { register, handleSubmit, formState, reset } = useForm<SimpleImportModelConfig>({
1817
defaultValues: {
@@ -22,34 +21,22 @@ export const InstallModelForm = () => {
2221
mode: 'onChange',
2322
});
2423

24+
const resetForm = useCallback(() => reset(undefined, { keepValues: true }), [reset]);
25+
2526
const onSubmit = useCallback<SubmitHandler<SimpleImportModelConfig>>(
2627
(values) => {
2728
if (!values?.location) {
2829
return;
2930
}
3031

31-
installModel({ source: values.location, inplace: values.inplace })
32-
.unwrap()
33-
.then((_) => {
34-
toast({
35-
id: ToastID.MODEL_INSTALL_QUEUED,
36-
title: t('toast.modelAddedSimple'),
37-
status: 'success',
38-
});
39-
reset(undefined, { keepValues: true });
40-
})
41-
.catch((error) => {
42-
reset(undefined, { keepValues: true });
43-
if (error) {
44-
toast({
45-
id: ToastID.MODEL_INSTALL_QUEUE_FAILED,
46-
title: `${error.data.detail} `,
47-
status: 'error',
48-
});
49-
}
50-
});
32+
installModel({
33+
source: values.location,
34+
inplace: values.inplace,
35+
onSuccess: resetForm,
36+
onError: resetForm,
37+
});
5138
},
52-
[reset, installModel]
39+
[installModel, resetForm]
5340
);
5441

5542
return (

‎invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/ScanFolder/ScanFolderResults.tsx

Lines changed: 7 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,12 @@ import {
1212
InputRightElement,
1313
} from '@invoke-ai/ui-library';
1414
import ScrollableContent from 'common/components/OverlayScrollbars/ScrollableContent';
15-
import { toast, ToastID } from 'features/toast/toast';
15+
import { useInstallModel } from 'features/modelManagerV2/hooks/useInstallModel';
1616
import type { ChangeEvent, ChangeEventHandler } from 'react';
1717
import { useCallback, useMemo, useState } from 'react';
1818
import { useTranslation } from 'react-i18next';
1919
import { PiXBold } from 'react-icons/pi';
20-
import { type ScanFolderResponse, useInstallModelMutation } from 'services/api/endpoints/models';
20+
import type { ScanFolderResponse } from 'services/api/endpoints/models';
2121

2222
import { ScanModelResultItem } from './ScanFolderResultItem';
2323

@@ -29,7 +29,7 @@ export const ScanModelsResults = ({ results }: ScanModelResultsProps) => {
2929
const { t } = useTranslation();
3030
const [searchTerm, setSearchTerm] = useState('');
3131
const [inplace, setInplace] = useState(true);
32-
const [installModel] = useInstallModelMutation();
32+
const [installModel] = useInstallModel();
3333

3434
const filteredResults = useMemo(() => {
3535
return results.filter((result) => {
@@ -55,49 +55,15 @@ export const ScanModelsResults = ({ results }: ScanModelResultsProps) => {
5555
if (result.is_installed) {
5656
continue;
5757
}
58-
installModel({ source: result.path, inplace })
59-
.unwrap()
60-
.then((_) => {
61-
toast({
62-
id: ToastID.MODEL_INSTALL_QUEUED,
63-
title: t('toast.modelAddedSimple'),
64-
status: 'success',
65-
});
66-
})
67-
.catch((error) => {
68-
if (error) {
69-
toast({
70-
id: ToastID.MODEL_INSTALL_QUEUE_FAILED,
71-
title: `${error.data.detail} `,
72-
status: 'error',
73-
});
74-
}
75-
});
58+
installModel({ source: result.path, inplace });
7659
}
77-
}, [filteredResults, installModel, inplace, t]);
60+
}, [filteredResults, installModel, inplace]);
7861

7962
const handleInstallOne = useCallback(
8063
(source: string) => {
81-
installModel({ source, inplace })
82-
.unwrap()
83-
.then((_) => {
84-
toast({
85-
id: ToastID.MODEL_INSTALL_QUEUED,
86-
title: t('toast.modelAddedSimple'),
87-
status: 'success',
88-
});
89-
})
90-
.catch((error) => {
91-
if (error) {
92-
toast({
93-
id: ToastID.MODEL_INSTALL_QUEUE_FAILED,
94-
title: `${error.data.detail} `,
95-
status: 'error',
96-
});
97-
}
98-
});
64+
installModel({ source, inplace });
9965
},
100-
[installModel, inplace, t]
66+
[installModel, inplace]
10167
);
10268

10369
return (

0 commit comments

Comments
 (0)