From df858eb3f969ab1a072b8962f683e96568dc8c7a Mon Sep 17 00:00:00 2001
From: blessedcoolant <54517381+blessedcoolant@users.noreply.github.com>
Date: Mon, 26 Dec 2022 23:47:49 +1300
Subject: [PATCH] Add Edit Model Functionality
---
frontend/public/locales/modelmanager/en.json | 3 +-
frontend/src/app/invokeai.d.ts | 16 +-
.../components/ModelManager/ModelEdit.tsx | 290 ++++++++++++++++--
3 files changed, 281 insertions(+), 28 deletions(-)
diff --git a/frontend/public/locales/modelmanager/en.json b/frontend/public/locales/modelmanager/en.json
index 48b2821ec7..aa824590af 100644
--- a/frontend/public/locales/modelmanager/en.json
+++ b/frontend/public/locales/modelmanager/en.json
@@ -11,7 +11,7 @@
"manual": "Manual",
"name": "Name",
"nameValidationMsg": "Enter a name for your model",
- "description": "description",
+ "description": "Description",
"descriptionValidationMsg": "Add a description for your model",
"config": "Config",
"configValidationMsg": "Path to the config file of your model.",
@@ -24,6 +24,7 @@
"height": "Height",
"heightValidationMsg": "Default height of your model.",
"addModel": "Add Model",
+ "updateModel": "Update Model",
"availableModels": "Available Models",
"search": "Search",
"load": "Load",
diff --git a/frontend/src/app/invokeai.d.ts b/frontend/src/app/invokeai.d.ts
index af292058c2..d5ee6adbbc 100644
--- a/frontend/src/app/invokeai.d.ts
+++ b/frontend/src/app/invokeai.d.ts
@@ -180,14 +180,14 @@ export declare type FoundModel = {
};
export declare type InvokeModelConfigProps = {
- name: string;
- description: string;
- config: string;
- weights: string;
- vae: string;
- width: number;
- height: number;
- default: boolean;
+ name: string | undefined;
+ description: string | undefined;
+ config: string | undefined;
+ weights: string | undefined;
+ vae: string | undefined;
+ width: number | undefined;
+ height: number | undefined;
+ default: boolean | undefined;
};
/**
diff --git a/frontend/src/features/system/components/ModelManager/ModelEdit.tsx b/frontend/src/features/system/components/ModelManager/ModelEdit.tsx
index 4738b0a2a7..11cbf8e2fe 100644
--- a/frontend/src/features/system/components/ModelManager/ModelEdit.tsx
+++ b/frontend/src/features/system/components/ModelManager/ModelEdit.tsx
@@ -1,11 +1,30 @@
-import { createSelector, Dictionary } from '@reduxjs/toolkit';
+import { createSelector } from '@reduxjs/toolkit';
import type { RootState } from 'app/store';
-import { useAppSelector } from 'app/storeHooks';
+import { useAppDispatch, useAppSelector } from 'app/storeHooks';
import { systemSelector } from 'features/system/store/systemSelectors';
import React, { useEffect, useState } from 'react';
import _ from 'lodash';
-import { Model } from 'app/invokeai';
-import { Flex, Spacer, Text } from '@chakra-ui/react';
+import type { InvokeModelConfigProps } from 'app/invokeai';
+import {
+ Button,
+ Flex,
+ FormControl,
+ FormErrorMessage,
+ FormHelperText,
+ FormLabel,
+ Input,
+ NumberDecrementStepper,
+ NumberIncrementStepper,
+ NumberInput,
+ NumberInputField,
+ NumberInputStepper,
+ Text,
+ VStack,
+} from '@chakra-ui/react';
+import { Field, Formik } from 'formik';
+import type { FieldInputProps, FormikProps } from 'formik';
+import { useTranslation } from 'react-i18next';
+import { addNewModel } from 'app/socketio/actions';
const selector = createSelector(
[systemSelector],
@@ -23,38 +42,271 @@ const selector = createSelector(
}
);
+const MIN_MODEL_SIZE = 64;
+const MAX_MODEL_SIZE = 2048;
+
export default function ModelEdit() {
const { openModel, model_list } = useAppSelector(selector);
- const [openedModel, setOpenedModel] = useState<Model>();
+ const isProcessing = useAppSelector(
+ (state: RootState) => state.system.isProcessing
+ );
+
+ const dispatch = useAppDispatch();
+
+ const { t } = useTranslation();
+
+ const [editModelFormValues, setEditModelFormValues] =
+ useState<InvokeModelConfigProps>({
+ name: '',
+ description: '',
+ config: 'configs/stable-diffusion/v1-inference.yaml',
+ weights: '',
+ vae: '',
+ width: 512,
+ height: 512,
+ default: false,
+ });
useEffect(() => {
if (openModel) {
const retrievedModel = _.pickBy(model_list, (val, key) => {
return _.isEqual(key, openModel);
});
- setOpenedModel(retrievedModel[openModel]);
+ setEditModelFormValues({
+ name: openModel,
+ description: retrievedModel[openModel]?.description,
+ config: retrievedModel[openModel]?.config,
+ weights: retrievedModel[openModel]?.weights,
+ vae: retrievedModel[openModel]?.vae,
+ width: retrievedModel[openModel]?.width,
+ height: retrievedModel[openModel]?.height,
+ default: retrievedModel[openModel]?.default,
+ });
}
}, [model_list, openModel]);
- console.log(openedModel);
+ const editModelFormSubmitHandler = (values: InvokeModelConfigProps) => {
+ dispatch(addNewModel(values));
+ };
- return (
- <Flex flexDirection="column" rowGap="1rem">
- <Flex columnGap="1rem" alignItems="center">
+ return openModel ? (
+ <Flex flexDirection="column" rowGap="1rem" width="100%">
+ <Flex alignItems="center">
<Text fontSize="lg" fontWeight="bold">
{openModel}
</Text>
- <Text>{openedModel?.status}</Text>
</Flex>
- <Flex flexDirection="column">
- <Text>{openedModel?.config}</Text>
- <Text>{openedModel?.description}</Text>
- <Text>{openedModel?.height}</Text>
- <Text>{openedModel?.width}</Text>
- <Text>{openedModel?.default}</Text>
- <Text>{openedModel?.vae}</Text>
- <Text>{openedModel?.weights}</Text>
+ <Flex
+ flexDirection="column"
+ maxHeight={window.innerHeight - 270}
+ overflowY="scroll"
+ paddingRight="2rem"
+ >
+ <Formik
+ enableReinitialize={true}
+ initialValues={editModelFormValues}
+ onSubmit={editModelFormSubmitHandler}
+ >
+ {({ handleSubmit, errors, touched }) => (
+ <form onSubmit={handleSubmit}>
+ <VStack rowGap={'0.5rem'} alignItems="start">
+ {/* Description */}
+ <FormControl
+ isInvalid={!!errors.description && touched.description}
+ isRequired
+ >
+ <FormLabel htmlFor="description">
+ {t('modelmanager:description')}
+ </FormLabel>
+ <VStack alignItems={'start'}>
+ <Field
+ as={Input}
+ id="description"
+ name="description"
+ type="text"
+ />
+ {!!errors.description && touched.description ? (
+ <FormErrorMessage>{errors.description}</FormErrorMessage>
+ ) : (
+ <FormHelperText margin={0}>
+ {t('modelmanager:descriptionValidationMsg')}
+ </FormHelperText>
+ )}
+ </VStack>
+ </FormControl>
+
+ {/* Config */}
+ <FormControl
+ isInvalid={!!errors.config && touched.config}
+ isRequired
+ >
+ <FormLabel htmlFor="config">
+ {t('modelmanager:config')}
+ </FormLabel>
+ <VStack alignItems={'start'}>
+ <Field as={Input} id="config" name="config" type="text" />
+ {!!errors.config && touched.config ? (
+ <FormErrorMessage>{errors.config}</FormErrorMessage>
+ ) : (
+ <FormHelperText margin={0}>
+ {t('modelmanager:configValidationMsg')}
+ </FormHelperText>
+ )}
+ </VStack>
+ </FormControl>
+
+ {/* Weights */}
+ <FormControl
+ isInvalid={!!errors.weights && touched.weights}
+ isRequired
+ >
+ <FormLabel htmlFor="config">
+ {t('modelmanager:modelLocation')}
+ </FormLabel>
+ <VStack alignItems={'start'}>
+ <Field as={Input} id="weights" name="weights" type="text" />
+ {!!errors.weights && touched.weights ? (
+ <FormErrorMessage>{errors.weights}</FormErrorMessage>
+ ) : (
+ <FormHelperText margin={0}>
+ {t('modelmanager:modelLocationValidationMsg')}
+ </FormHelperText>
+ )}
+ </VStack>
+ </FormControl>
+
+ {/* VAE */}
+ <FormControl isInvalid={!!errors.vae && touched.vae}>
+ <FormLabel htmlFor="vae">
+ {t('modelmanager:vaeLocation')}
+ </FormLabel>
+ <VStack alignItems={'start'}>
+ <Field as={Input} id="vae" name="vae" type="text" />
+ {!!errors.vae && touched.vae ? (
+ <FormErrorMessage>{errors.vae}</FormErrorMessage>
+ ) : (
+ <FormHelperText margin={0}>
+ {t('modelmanager:vaeLocationValidationMsg')}
+ </FormHelperText>
+ )}
+ </VStack>
+ </FormControl>
+
+ <VStack width={'100%'}>
+ {/* Width */}
+ <FormControl isInvalid={!!errors.width && touched.width}>
+ <FormLabel htmlFor="width">
+ {t('modelmanager:width')}
+ </FormLabel>
+ <VStack alignItems={'start'}>
+ <Field id="width" name="width">
+ {({
+ field,
+ form,
+ }: {
+ field: FieldInputProps<number>;
+ form: FormikProps<InvokeModelConfigProps>;
+ }) => (
+ <NumberInput
+ {...field}
+ id="width"
+ name="width"
+ min={MIN_MODEL_SIZE}
+ max={MAX_MODEL_SIZE}
+ step={64}
+ onChange={(value) =>
+ form.setFieldValue(field.name, Number(value))
+ }
+ >
+ <NumberInputField />
+ <NumberInputStepper>
+ <NumberIncrementStepper />
+ <NumberDecrementStepper />
+ </NumberInputStepper>
+ </NumberInput>
+ )}
+ </Field>
+
+ {!!errors.width && touched.width ? (
+ <FormErrorMessage>{errors.width}</FormErrorMessage>
+ ) : (
+ <FormHelperText margin={0}>
+ {t('modelmanager:widthValidationMsg')}
+ </FormHelperText>
+ )}
+ </VStack>
+ </FormControl>
+
+ {/* Height */}
+ <FormControl isInvalid={!!errors.height && touched.height}>
+ <FormLabel htmlFor="height">
+ {t('modelmanager:height')}
+ </FormLabel>
+ <VStack alignItems={'start'}>
+ <Field id="height" name="height">
+ {({
+ field,
+ form,
+ }: {
+ field: FieldInputProps<number>;
+ form: FormikProps<InvokeModelConfigProps>;
+ }) => (
+ <NumberInput
+ {...field}
+ id="height"
+ name="height"
+ min={MIN_MODEL_SIZE}
+ max={MAX_MODEL_SIZE}
+ step={64}
+ onChange={(value) =>
+ form.setFieldValue(field.name, Number(value))
+ }
+ >
+ <NumberInputField />
+ <NumberInputStepper>
+ <NumberIncrementStepper />
+ <NumberDecrementStepper />
+ </NumberInputStepper>
+ </NumberInput>
+ )}
+ </Field>
+
+ {!!errors.height && touched.height ? (
+ <FormErrorMessage>{errors.height}</FormErrorMessage>
+ ) : (
+ <FormHelperText margin={0}>
+ {t('modelmanager:heightValidationMsg')}
+ </FormHelperText>
+ )}
+ </VStack>
+ </FormControl>
+ </VStack>
+
+ <Button
+ type="submit"
+ className="modal-close-btn"
+ isLoading={isProcessing}
+ >
+ {t('modelmanager:updateModel')}
+ </Button>
+ </VStack>
+ </form>
+ )}
+ </Formik>
</Flex>
</Flex>
+ ) : (
+ <Flex
+ width="100%"
+ height="250px"
+ justifyContent="center"
+ alignItems="center"
+ backgroundColor="var(--background-color)"
+ borderRadius="0.5rem"
+ >
+ <Text fontWeight="bold" color="var(--subtext-color-bright)">
+ Pick A Model To Edit
+ </Text>
+ </Flex>
);
}