Skip to content

Commit fbf7e2e

Browse files
magnusvirimagnusviri
committed
Scaffolding for invoke-ai#6179, upload model via model manager
1 parent 9066dc1 commit fbf7e2e

3 files changed

Lines changed: 283 additions & 0 deletions

File tree

‎invokeai/frontend/web/src/common/hooks/useModelUploadButton.tsx‎

Lines changed: 249 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,249 @@
1+
import type { ButtonProps, IconButtonProps, SystemStyleObject } from '@invoke-ai/ui-library';
2+
import { Button, IconButton } from '@invoke-ai/ui-library';
3+
import { logger } from 'app/logging/logger';
4+
import { useAppSelector } from 'app/store/storeHooks';
5+
import { selectAutoAddBoardId } from 'features/gallery/store/gallerySelectors';
6+
import { selectIsClientSideUploadEnabled } from 'features/system/store/configSlice';
7+
import { toast } from 'features/toast/toast';
8+
import { memo, useCallback } from 'react';
9+
import type { FileRejection } from 'react-dropzone';
10+
import { useDropzone } from 'react-dropzone';
11+
import { useTranslation } from 'react-i18next';
12+
import { PiUploadBold } from 'react-icons/pi';
13+
import { uploadImages, useUploadImageMutation } from 'services/api/endpoints/images';
14+
import type { ImageDTO } from 'services/api/types';
15+
import { assert } from 'tsafe';
16+
import type { SetOptional } from 'type-fest';
17+
18+
import { useClientSideUpload } from './useClientSideUpload';
19+
type useModelUploadButtonArgs =
20+
| {
21+
isDisabled?: boolean;
22+
allowMultiple: false;
23+
onUpload?: (imageDTO: ImageDTO) => void;
24+
}
25+
| {
26+
isDisabled?: boolean;
27+
allowMultiple: true;
28+
onUpload?: (imageDTOs: ImageDTO[]) => void;
29+
};
30+
31+
const log = logger('gallery');
32+
33+
/**
34+
* Provides model uploader functionality to any component.
35+
*
36+
* @example
37+
* const { getUploadButtonProps, getUploadInputProps, openUploader } = useModelUploadButton({
38+
* postUploadAction: {
39+
* type: 'SET_CONTROL_ADAPTER_IMAGE',
40+
* controlNetId: '12345',
41+
* },
42+
* isDisabled: getIsUploadDisabled(),
43+
* });
44+
*
45+
* // open the uploaded directly
46+
* const handleSomething = () => { openUploader() }
47+
*
48+
* // in the render function
49+
* <Button {...getUploadButtonProps()} /> // will open the file dialog on click
50+
* <input {...getUploadInputProps()} /> // hidden, handles native upload functionality
51+
*/
52+
export const useModelUploadButton = ({ onUpload, isDisabled, allowMultiple }: useModelUploadButtonArgs) => {
53+
const autoAddBoardId = useAppSelector(selectAutoAddBoardId);
54+
const isClientSideUploadEnabled = useAppSelector(selectIsClientSideUploadEnabled);
55+
const [uploadImage, request] = useUploadImageMutation();
56+
const clientSideUpload = useClientSideUpload();
57+
const { t } = useTranslation();
58+
59+
const onDropAccepted = useCallback(
60+
async (files: File[]) => {
61+
try {
62+
if (!allowMultiple) {
63+
if (files.length > 1) {
64+
log.warn('Multiple files dropped but only one allowed');
65+
return;
66+
}
67+
if (files.length === 0) {
68+
// Should never happen
69+
log.warn('No files dropped');
70+
return;
71+
}
72+
const file = files[0];
73+
assert(file !== undefined); // should never happen
74+
const imageDTO = await uploadImage({
75+
file,
76+
image_category: 'user',
77+
is_intermediate: false,
78+
board_id: autoAddBoardId === 'none' ? undefined : autoAddBoardId,
79+
silent: true,
80+
}).unwrap();
81+
if (onUpload) {
82+
onUpload(imageDTO);
83+
}
84+
} else {
85+
let imageDTOs: ImageDTO[] = [];
86+
if (isClientSideUploadEnabled && files.length > 1) {
87+
imageDTOs = await Promise.all(files.map((file, i) => clientSideUpload(file, i)));
88+
} else {
89+
imageDTOs = await uploadImages(
90+
files.map((file, i) => ({
91+
file,
92+
image_category: 'user',
93+
is_intermediate: false,
94+
board_id: autoAddBoardId === 'none' ? undefined : autoAddBoardId,
95+
silent: false,
96+
isFirstUploadOfBatch: i === 0,
97+
}))
98+
);
99+
}
100+
if (onUpload) {
101+
onUpload(imageDTOs);
102+
}
103+
}
104+
} catch (error) {
105+
toast({
106+
id: 'UPLOAD_FAILED',
107+
title: t('toast.imageUploadFailed'),
108+
status: 'error',
109+
});
110+
}
111+
},
112+
[allowMultiple, autoAddBoardId, onUpload, uploadImage, isClientSideUploadEnabled, clientSideUpload, t]
113+
);
114+
115+
const onDropRejected = useCallback(
116+
(fileRejections: FileRejection[]) => {
117+
if (fileRejections.length > 0) {
118+
const errors = fileRejections.map((rejection) => ({
119+
errors: rejection.errors.map(({ message }) => message),
120+
file: rejection.file.path,
121+
}));
122+
log.error({ errors }, 'Invalid upload');
123+
const description = t('toast.uploadFailedInvalidUploadDesc');
124+
125+
toast({
126+
id: 'UPLOAD_FAILED',
127+
title: t('toast.uploadFailed'),
128+
description,
129+
status: 'error',
130+
});
131+
132+
return;
133+
}
134+
},
135+
[t]
136+
);
137+
138+
const {
139+
getRootProps: getUploadButtonProps,
140+
getInputProps: getUploadInputProps,
141+
open: openUploader,
142+
} = useDropzone({
143+
accept: {
144+
'image/png': ['.png'],
145+
'image/jpeg': ['.jpg', '.jpeg', '.png'],
146+
'image/webp': ['.webp'],
147+
},
148+
onDropAccepted,
149+
onDropRejected,
150+
disabled: isDisabled,
151+
noDrag: true,
152+
multiple: allowMultiple,
153+
});
154+
155+
return { getUploadButtonProps, getUploadInputProps, openUploader, request };
156+
};
157+
158+
const sx = {
159+
'&[data-error=true]': {
160+
borderColor: 'error.500',
161+
borderStyle: 'solid',
162+
borderWidth: 1,
163+
},
164+
} satisfies SystemStyleObject;
165+
166+
export const UploadModelIconButton = memo(
167+
({
168+
isDisabled = false,
169+
onUpload,
170+
isError = false,
171+
...rest
172+
}: {
173+
onUpload?: (imageDTO: ImageDTO) => void;
174+
isError?: boolean;
175+
} & SetOptional<IconButtonProps, 'aria-label'>) => {
176+
const uploadApi = useModelUploadButton({ isDisabled, allowMultiple: false, onUpload });
177+
return (
178+
<>
179+
<IconButton
180+
aria-label="Upload model"
181+
variant="outline"
182+
sx={sx}
183+
data-error={isError}
184+
icon={<PiUploadBold />}
185+
isLoading={uploadApi.request.isLoading}
186+
{...rest}
187+
{...uploadApi.getUploadButtonProps()}
188+
/>
189+
<input {...uploadApi.getUploadInputProps()} />
190+
</>
191+
);
192+
}
193+
);
194+
UploadModelIconButton.displayName = 'UploadModelIconButton';
195+
196+
type UploadModelButtonProps = {
197+
onUpload?: (imageDTO: ImageDTO) => void;
198+
isError?: boolean;
199+
} & ButtonProps;
200+
201+
const UploadModelButton = memo((props: UploadModelButtonProps) => {
202+
const { children, isDisabled = false, onUpload, isError = false, ...rest } = props;
203+
const uploadApi = useModelUploadButton({ isDisabled, allowMultiple: false, onUpload });
204+
return (
205+
<>
206+
<Button
207+
aria-label="Upload model"
208+
variant="outline"
209+
sx={sx}
210+
data-error={isError}
211+
rightIcon={<PiUploadBold />}
212+
isLoading={uploadApi.request.isLoading}
213+
{...rest}
214+
{...uploadApi.getUploadButtonProps()}
215+
>
216+
{children ?? 'Upload'}
217+
</Button>
218+
<input {...uploadApi.getUploadInputProps()} />
219+
</>
220+
);
221+
});
222+
UploadModelButton.displayName = 'UploadModelButton';
223+
224+
export const UploadMultipleModelButton = ({
225+
isDisabled = false,
226+
onUpload,
227+
isError = false,
228+
...rest
229+
}: {
230+
onUpload?: (imageDTOs: ImageDTO[]) => void;
231+
isError?: boolean;
232+
} & SetOptional<IconButtonProps, 'aria-label'>) => {
233+
const uploadApi = useModelUploadButton({ isDisabled, allowMultiple: true, onUpload });
234+
return (
235+
<>
236+
<IconButton
237+
aria-label="Upload model"
238+
variant="outline"
239+
sx={sx}
240+
data-error={isError}
241+
icon={<PiUploadBold />}
242+
isLoading={uploadApi.request.isLoading}
243+
{...rest}
244+
{...uploadApi.getUploadButtonProps()}
245+
/>
246+
<input {...uploadApi.getUploadInputProps()} />
247+
</>
248+
);
249+
};

‎invokeai/frontend/web/src/features/modelManagerV2/components/ModelUploadButton.tsx‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
import { IconButton } from '@invoke-ai/ui-library';
2+
import { useModelUploadButton } from 'common/hooks/useModelUploadButton';
3+
import { t } from 'i18next';
4+
import { memo } from 'react';
5+
import { PiUploadBold } from 'react-icons/pi';
6+
7+
const UPLOAD_OPTIONS: Parameters<typeof useModelUploadButton>[0] = { allowMultiple: true };
8+
9+
export const ModelUploadButton = memo(() => {
10+
const uploadApi = useModelUploadButton(UPLOAD_OPTIONS);
11+
return (
12+
<>
13+
<IconButton
14+
size="sm"
15+
alignSelf="stretch"
16+
variant="link"
17+
aria-label={t('accessibility.uploadImages')}
18+
tooltip={t('accessibility.uploadImages')}
19+
icon={<PiUploadBold />}
20+
{...uploadApi.getUploadButtonProps()}
21+
/>
22+
<input {...uploadApi.getUploadInputProps()} />
23+
</>
24+
);
25+
});
26+
ModelUploadButton.displayName = 'ModelUploadButton';

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import { Button, Checkbox, Flex, FormControl, FormHelperText, FormLabel, Input } from '@invoke-ai/ui-library';
22
import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
3+
import { ModelUploadButton } from 'features/modelManagerV2/components/ModelUploadButton';
34
import { useInstallModel } from 'features/modelManagerV2/hooks/useInstallModel';
45
import {
56
selectShouldInstallInPlace,
@@ -83,6 +84,13 @@ export const InstallModelForm = memo(() => {
8384
</Flex>
8485
</FormControl>
8586
</Flex>
87+
88+
<FormControl orientation="vertical">
89+
<FormLabel>Choose File</FormLabel>
90+
<Flex alignItems="center" gap={3} w="full">
91+
<ModelUploadButton />
92+
</Flex>
93+
</FormControl>
8694
</form>
8795
);
8896
});

0 commit comments

Comments
 (0)