moved upload image field and added delete image functionality

This commit is contained in:
Jennifer Player 2024-03-06 09:58:39 -05:00 committed by Kent Keirsey
parent 8a68355926
commit 239b1e8cc7
6 changed files with 105 additions and 81 deletions

View File

@ -292,6 +292,7 @@ async def get_model_image(
"""Gets a full-resolution image file"""
try:
# still need to handle this gracefully when path doesnt exist instead of throwing error
path = ApiDependencies.invoker.services.model_images.get_path(key + ".png")
if not path:
@ -308,22 +309,6 @@ async def get_model_image(
except Exception:
raise HTTPException(status_code=404)
# async def get_model_image(
# key: Annotated[str, Path(description="Unique key of model")],
# ) -> Optional[str]:
# model_images = ApiDependencies.invoker.services.model_images
# try:
# url = model_images.get_url(key)
# if not url:
# return None
# return url
# except Exception:
# raise HTTPException(status_code=404)
@model_manager_router.patch(
"/i/{key}/image",
operation_id="update_model_image",
@ -390,6 +375,28 @@ async def delete_model(
logger.error(str(e))
raise HTTPException(status_code=404, detail=str(e))
@model_manager_router.delete(
"/i/{key}/image",
operation_id="delete_model_image",
responses={
204: {"description": "Model image deleted successfully"},
404: {"description": "Model image not found"},
},
status_code=204,
)
async def delete_model_image(
key: str = Path(description="Unique key of model image to remove from model_images directory."),
) -> None:
logger = ApiDependencies.invoker.services.logger
model_images = ApiDependencies.invoker.services.model_images
try:
model_images.delete(key)
logger.info(f"Deleted model image: {key}")
return
except UnknownModelException as e:
logger.error(str(e))
raise HTTPException(status_code=404, detail=str(e))
# @model_manager_router.post(
# "/i/",

View File

@ -16,11 +16,6 @@ class ModelImagesBase(ABC):
"""Gets the internal path to a model image."""
pass
@abstractmethod
def get_url(self, model_key: str) -> str:
"""Gets the url for a model image."""
pass
@abstractmethod
def save(
self,

View File

@ -1,25 +1,28 @@
import { Box, IconButton, Image } from '@invoke-ai/ui-library';
import { typedMemo } from 'common/util/typedMemo';
import { useCallback, useState } from 'react';
import type { Control } from 'react-hook-form';
import { useController, useWatch } from 'react-hook-form';
import { useAppDispatch } from 'app/store/storeHooks';
import { Button } from '@invoke-ai/ui-library';
import { useDropzone } from 'react-dropzone';
import { PiArrowCounterClockwiseBold, PiUploadSimpleBold } from 'react-icons/pi';
import { UpdateModelArg, buildModelsUrl } from 'services/api/endpoints/models';
import { buildModelsUrl, useUpdateModelImageMutation, useDeleteModelImageMutation } from 'services/api/endpoints/models';
import { useTranslation } from 'react-i18next';
import { addToast } from 'features/system/store/systemSlice';
import { makeToast } from 'features/system/util/makeToast';
type Props = {
control: Control<UpdateModelArg['body']>;
model_key: string;
};
const ModelImageUpload = ({ control }: Props) => {
const key = useWatch({ control, name: 'key' });
const [image, setImage] = useState<string | undefined>(buildModelsUrl(`i/${key}/image`));
const { field } = useController({ control, name: 'image' });
const ModelImageUpload = ({ model_key }: Props) => {
const dispatch = useAppDispatch();
const [image, setImage] = useState<string | undefined>(buildModelsUrl(`i/${model_key}/image`));
const { t } = useTranslation();
const [updateModelImage] = useUpdateModelImageMutation();
const [deleteModelImage] = useDeleteModelImageMutation();
const onDropAccepted = useCallback(
(files: File[]) => {
const file = files[0];
@ -28,16 +31,59 @@ const ModelImageUpload = ({ control }: Props) => {
return;
}
field.onChange(file);
setImage(URL.createObjectURL(file));
updateModelImage({ key: model_key, image: image })
.unwrap()
.then(() => {
dispatch(
addToast(
makeToast({
title: t('modelManager.modelImageUpdated'),
status: 'success',
})
)
);
})
.catch((_) => {
dispatch(
addToast(
makeToast({
title: t('modelManager.modelImageUpdateFailed'),
status: 'error',
})
)
);
});
},
[field]
[]
);
const handleResetImage = useCallback(() => {
field.onChange(undefined);
setImage(undefined);
}, [field]);
deleteModelImage(model_key)
.unwrap()
.then(() => {
dispatch(
addToast(
makeToast({
title: t('modelManager.modelImageDeleted'),
status: 'success',
})
)
);
})
.catch((_) => {
dispatch(
addToast(
makeToast({
title: t('modelManager.modelImageDeleteFailed'),
status: 'error',
})
)
);
});
}, []);
const { getInputProps, getRootProps } = useDropzone({
accept: { 'image/png': ['.png'], 'image/jpeg': ['.jpg', '.jpeg', '.png'] },
@ -50,8 +96,6 @@ const ModelImageUpload = ({ control }: Props) => {
return (
<Box
position="relative"
_hover={{ filter: 'brightness(50%)' }}
transition="filter ease 0.2s"
>
<Image
onError={() => setImage(undefined)}
@ -66,11 +110,12 @@ const ModelImageUpload = ({ control }: Props) => {
top="1"
right="1"
onClick={handleResetImage}
aria-label={t('modelManager.resetImage')}
tooltip={t('modelManager.resetImage')}
aria-label={t('modelManager.deleteModelImage')}
tooltip={t('modelManager.deleteModelImage')}
icon={<PiArrowCounterClockwiseBold size={16} />}
size="sm"
variant="link"
_hover={{ color: 'base.100' }}
/>
</Box>
);

View File

@ -8,6 +8,7 @@ import { ModelMetadata } from './Metadata/ModelMetadata';
import { ModelAttrView } from './ModelAttrView';
import { ModelEdit } from './ModelEdit';
import { ModelView } from './ModelView';
import ModelImageUpload from './Fields/ModelImageUpload';
export const Model = () => {
const { t } = useTranslation();
@ -25,6 +26,7 @@ export const Model = () => {
return (
<>
<Flex alignItems="center" justifyContent="space-between" gap="4" paddingRight="5">
<Flex flexDir="column" gap={1} p={2}>
<Heading as="h2" fontSize="lg">
{data.name}
@ -39,6 +41,8 @@ export const Model = () => {
<ModelAttrView label="Description" value={data.description} />
</Box>
</Flex>
<ModelImageUpload model_key={selectedModelKey || ''} />
</Flex>
<Tabs mt="4" h="100%">
<TabList>

View File

@ -20,14 +20,9 @@ import type { SubmitHandler } from 'react-hook-form';
import { useForm } from 'react-hook-form';
import { useTranslation } from 'react-i18next';
import type { UpdateModelArg } from 'services/api/endpoints/models';
import {
useGetModelConfigQuery,
useUpdateModelImageMutation,
useUpdateModelsMutation,
} from 'services/api/endpoints/models';
import { useGetModelConfigQuery, useUpdateModelMutation } from 'services/api/endpoints/models';
import BaseModelSelect from './Fields/BaseModelSelect';
import ModelImageUpload from './Fields/ModelImageUpload';
import ModelVariantSelect from './Fields/ModelVariantSelect';
import PredictionTypeSelect from './Fields/PredictionTypeSelect';
@ -36,8 +31,7 @@ export const ModelEdit = () => {
const selectedModelKey = useAppSelector((s) => s.modelmanagerV2.selectedModelKey);
const { data, isLoading } = useGetModelConfigQuery(selectedModelKey ?? skipToken);
const [updateModel, { isLoading: isSubmitting }] = useUpdateModelsMutation();
const [updateModelImage] = useUpdateModelImageMutation();
const [updateModel, { isLoading: isSubmitting }] = useUpdateModelMutation();
const { t } = useTranslation();
@ -60,11 +54,6 @@ export const ModelEdit = () => {
return;
}
// remove image from body
const image = values.image;
if (values.image) {
delete values.image;
}
const responseBody: UpdateModelArg = {
key: data.key,
body: values,
@ -95,33 +84,6 @@ export const ModelEdit = () => {
)
);
});
if (image) {
updateModelImage({ key: data.key, image: image })
.unwrap()
.then((payload) => {
reset(payload, { keepDefaultValues: true });
dispatch(setSelectedModelMode('view'));
dispatch(
addToast(
makeToast({
title: t('modelManager.modelUpdated'),
status: 'success',
})
)
);
})
.catch((_) => {
reset();
dispatch(
addToast(
makeToast({
title: t('modelManager.modelUpdateFailed'),
status: 'error',
})
)
);
});
}
},
[dispatch, data?.key, reset, t, updateModel]
);
@ -168,7 +130,6 @@ export const ModelEdit = () => {
<Flex flexDir="column" gap={3} mt="4">
<Flex gap="4" alignItems="center">
<ModelImageUpload control={control} />
<FormControl flexDir="column" alignItems="flex-start" gap={1}>
<FormLabel>{t('modelManager.description')}</FormLabel>
<Textarea fontSize="md" resize="none" {...register('description')} />

View File

@ -40,6 +40,7 @@ type DeleteModelArg = {
key: string;
};
type DeleteModelResponse = void;
type DeleteModelImageResponse = void;
type ConvertMainModelResponse =
paths['/api/v2/models/convert/{key}']['put']['responses']['200']['content']['application/json'];
@ -161,7 +162,6 @@ export const modelsApi = api.injectEndpoints({
body: formData,
};
},
invalidatesTags: ['Model'],
}),
installModel: build.mutation<InstallModelResponse, InstallModelArg>({
query: ({ source }) => {
@ -182,6 +182,17 @@ export const modelsApi = api.injectEndpoints({
},
invalidatesTags: ['Model'],
}),
deleteModelImage: build.mutation<DeleteModelImageResponse, string>({
query: (key) => {
return {
url: buildModelsUrl(`i/${key}/image`),
method: 'DELETE',
};
},
}),
getModelImage: build.query<string, string>({
query: (key) => buildModelsUrl(`i/${key}/image`)
}),
convertModel: build.mutation<ConvertMainModelResponse, string>({
query: (key) => {
return {
@ -349,6 +360,7 @@ export const {
useGetTextualInversionModelsQuery,
useGetVaeModelsQuery,
useDeleteModelsMutation,
useDeleteModelImageMutation,
useUpdateModelMutation,
useUpdateModelImageMutation,
useInstallModelMutation,