feat: Restore Model Convert Functionality

This commit is contained in:
blessedcoolant 2023-07-12 20:40:58 +12:00
parent 6238a53fdd
commit 2cedf6aed5
4 changed files with 61 additions and 27 deletions

View File

@ -409,7 +409,7 @@
"convertToDiffusersHelpText2": "This process will replace your Model Manager entry with the Diffusers version of the same model.",
"convertToDiffusersHelpText3": "Your checkpoint file on the disk will NOT be deleted or modified in anyway. You can add your checkpoint to the Model Manager again if you want to.",
"convertToDiffusersHelpText4": "This is a one time process only. It might take around 30s-60s depending on the specifications of your computer.",
"convertToDiffusersHelpText5": "Please make sure you have enough disk space. Models generally vary between 4GB-7GB in size.",
"convertToDiffusersHelpText5": "Please make sure you have enough disk space. Models generally vary between 2GB-7GB in size.",
"convertToDiffusersHelpText6": "Do you wish to convert this model?",
"convertToDiffusersSaveLocation": "Save Location",
"v1": "v1",
@ -420,6 +420,7 @@
"pathToCustomConfig": "Path To Custom Config",
"statusConverting": "Converting",
"modelConverted": "Model Converted",
"modelConversionFailed": "Model Conversion Failed",
"sameFolder": "Same folder",
"invokeRoot": "InvokeAI folder",
"custom": "Custom",

View File

@ -77,7 +77,7 @@ export default function DiffusersModelEdit(props: DiffusersModelEditProps) {
addToast(
makeToast({
title: t('modelManager.modelUpdateFailed'),
status: 'success',
status: 'error',
})
)
);

View File

@ -1,23 +1,17 @@
import {
Flex,
ListItem,
Radio,
RadioGroup,
Text,
Tooltip,
UnorderedList,
} from '@chakra-ui/react';
import { Flex, ListItem, Text, UnorderedList } from '@chakra-ui/react';
// import { convertToDiffusers } from 'app/socketio/actions';
import { makeToast } from 'app/components/Toaster';
import { useAppDispatch } from 'app/store/storeHooks';
import IAIAlertDialog from 'common/components/IAIAlertDialog';
import IAIButton from 'common/components/IAIButton';
import IAIInput from 'common/components/IAIInput';
import { addToast } from 'features/system/store/systemSlice';
import { useEffect, useState } from 'react';
import { useTranslation } from 'react-i18next';
import { CheckpointModel } from './CheckpointModelEdit';
import { useConvertMainModelMutation } from 'services/api/endpoints/models';
import { CheckpointModelConfig } from './CheckpointModelEdit';
interface ModelConvertProps {
model: CheckpointModel;
model: CheckpointModelConfig;
}
export default function ModelConvert(props: ModelConvertProps) {
@ -26,6 +20,9 @@ export default function ModelConvert(props: ModelConvertProps) {
const dispatch = useAppDispatch();
const { t } = useTranslation();
const [convertModel, { isLoading, error, data }] =
useConvertMainModelMutation();
const [saveLocation, setSaveLocation] = useState<string>('same');
const [customSaveLocation, setCustomSaveLocation] = useState<string>('');
@ -38,15 +35,33 @@ export default function ModelConvert(props: ModelConvertProps) {
};
const modelConvertHandler = () => {
const modelToConvert = {
model_name: model,
save_location: saveLocation,
custom_location:
saveLocation === 'custom' && customSaveLocation !== ''
? customSaveLocation
: null,
const responseBody = {
base_model: model.base_model,
model_name: model.name,
};
dispatch(convertToDiffusers(modelToConvert));
convertModel(responseBody);
if (error) {
dispatch(
addToast(
makeToast({
title: t('modelManager.modelConversionFailed'),
status: 'error',
})
)
);
}
if (data) {
dispatch(
addToast(
makeToast({
title: t('modelManager.modelConverted'),
status: 'success',
})
)
);
}
};
return (
@ -60,6 +75,7 @@ export default function ModelConvert(props: ModelConvertProps) {
size={'sm'}
aria-label={t('modelManager.convertToDiffusers')}
className=" modal-close-btn"
isLoading={isLoading}
>
🧨 {t('modelManager.convertToDiffusers')}
</IAIButton>
@ -77,7 +93,7 @@ export default function ModelConvert(props: ModelConvertProps) {
<Text>{t('modelManager.convertToDiffusersHelpText6')}</Text>
</Flex>
<Flex flexDir="column" gap={4}>
{/* <Flex flexDir="column" gap={4}>
<Flex marginTop={4} flexDir="column" gap={2}>
<Text fontWeight="600">
{t('modelManager.convertToDiffusersSaveLocation')}
@ -103,9 +119,9 @@ export default function ModelConvert(props: ModelConvertProps) {
</Radio>
</Flex>
</RadioGroup>
</Flex>
</Flex> */}
{saveLocation === 'custom' && (
{/* {saveLocation === 'custom' && (
<Flex flexDirection="column" rowGap={2}>
<Text fontWeight="500" fontSize="sm" variant="subtext">
{t('modelManager.customSaveLocation')}
@ -119,8 +135,7 @@ export default function ModelConvert(props: ModelConvertProps) {
width="full"
/>
</Flex>
)}
</Flex>
)} */}
</IAIAlertDialog>
);
}

View File

@ -44,6 +44,11 @@ type DeleteMainModelQuery = {
model_name: string;
};
type ConvertMainModelQuery = {
base_model: BaseModelType;
model_name: string;
};
const mainModelsAdapter = createEntityAdapter<MainModelConfigEntity>({
sortComparer: (a, b) => a.name.localeCompare(b.name),
});
@ -138,6 +143,18 @@ export const modelsApi = api.injectEndpoints({
},
invalidatesTags: ['MainModel'],
}),
convertMainModel: build.mutation<
EntityState<MainModelConfigEntity>,
ConvertMainModelQuery
>({
query: ({ base_model, model_name }) => {
return {
url: `models/convert/${base_model}/main/${model_name}`,
method: 'PUT',
};
},
invalidatesTags: ['MainModel'],
}),
getLoRAModels: build.query<EntityState<LoRAModelConfigEntity>, void>({
query: () => ({ url: 'models/', params: { model_type: 'lora' } }),
providesTags: (result, error, arg) => {
@ -283,4 +300,5 @@ export const {
useGetVaeModelsQuery,
useUpdateMainModelsMutation,
useDeleteMainModelsMutation,
useConvertMainModelMutation,
} = modelsApi;