feat: Restore Update Model functionality

This commit is contained in:
blessedcoolant
2023-07-12 16:13:49 +12:00
parent af239fa122
commit afb46564e8
4 changed files with 108 additions and 33 deletions

View File

@ -2,6 +2,7 @@ import { EntityState, createEntityAdapter } from '@reduxjs/toolkit';
import { cloneDeep } from 'lodash-es';
import {
AnyModelConfig,
BaseModelType,
ControlNetModelConfig,
LoRAModelConfig,
MainModelConfig,
@ -32,6 +33,12 @@ type AnyModelConfigEntity =
| TextualInversionModelConfigEntity
| VaeModelConfigEntity;
type UpdateMainModelQuery = {
base_model: BaseModelType;
model_name: string;
body: MainModelConfig;
};
const mainModelsAdapter = createEntityAdapter<MainModelConfigEntity>({
sortComparer: (a, b) => a.name.localeCompare(b.name),
});
@ -101,6 +108,19 @@ export const modelsApi = api.injectEndpoints({
);
},
}),
updateMainModels: build.mutation<
EntityState<MainModelConfigEntity>,
UpdateMainModelQuery
>({
query: ({ base_model, model_name, body }) => {
return {
url: `models/${base_model}/main/${model_name}`,
method: 'PATCH',
body: body,
};
},
invalidatesTags: ['MainModel'],
}),
getLoRAModels: build.query<EntityState<LoRAModelConfigEntity>, void>({
query: () => ({ url: 'models/', params: { model_type: 'lora' } }),
providesTags: (result, error, arg) => {
@ -244,4 +264,5 @@ export const {
useGetLoRAModelsQuery,
useGetTextualInversionModelsQuery,
useGetVaeModelsQuery,
useUpdateMainModelsMutation,
} = modelsApi;