feat: Restore Model Merge functionality

This commit is contained in:
blessedcoolant
2023-07-12 22:43:06 +12:00
parent 683229e285
commit 3db1aa738c
5 changed files with 172 additions and 64 deletions

View File

@ -6,6 +6,7 @@ import {
ControlNetModelConfig,
LoRAModelConfig,
MainModelConfig,
MergeModelConfig,
TextualInversionModelConfig,
VaeModelConfig,
} from 'services/api/types';
@ -49,6 +50,11 @@ type ConvertMainModelQuery = {
model_name: string;
};
type MergeMainModelQuery = {
base_model: BaseModelType;
body: MergeModelConfig;
};
const mainModelsAdapter = createEntityAdapter<MainModelConfigEntity>({
sortComparer: (a, b) => a.name.localeCompare(b.name),
});
@ -143,7 +149,7 @@ export const modelsApi = api.injectEndpoints({
},
invalidatesTags: ['MainModel'],
}),
convertMainModel: build.mutation<
convertMainModels: build.mutation<
EntityState<MainModelConfigEntity>,
ConvertMainModelQuery
>({
@ -155,6 +161,19 @@ export const modelsApi = api.injectEndpoints({
},
invalidatesTags: ['MainModel'],
}),
mergeMainModels: build.mutation<
EntityState<MainModelConfigEntity>,
MergeMainModelQuery
>({
query: ({ base_model, body }) => {
return {
url: `models/merge/${base_model}`,
method: 'PUT',
body: body,
};
},
invalidatesTags: ['MainModel'],
}),
getLoRAModels: build.query<EntityState<LoRAModelConfigEntity>, void>({
query: () => ({ url: 'models/', params: { model_type: 'lora' } }),
providesTags: (result, error, arg) => {
@ -300,5 +319,6 @@ export const {
useGetVaeModelsQuery,
useUpdateMainModelsMutation,
useDeleteMainModelsMutation,
useConvertMainModelMutation,
useConvertMainModelsMutation,
useMergeMainModelsMutation,
} = modelsApi;