Add Edit Model Functionality

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",

frontend/src/app/invokeai.d.ts

@@ -180,14 +180,14 @@ export declare type FoundModel = {
 };
 
 export declare type InvokeModelConfigProps = {
-  name: string;
+  name: string | undefined;
-  description: string;
+  description: string | undefined;
-  config: string;
+  config: string | undefined;
-  weights: string;
+  weights: string | undefined;
-  vae: string;
+  vae: string | undefined;
-  width: number;
+  width: number | undefined;
-  height: number;
+  height: number | undefined;
-  default: boolean;
+  default: boolean | undefined;
 };
 
 /**

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 type { InvokeModelConfigProps } from 'app/invokeai';
-import { Flex, Spacer, Text } from '@chakra-ui/react';
+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 (
+  return openModel ? (
-    <Flex flexDirection="column" rowGap="1rem">
+    <Flex flexDirection="column" rowGap="1rem" width="100%">
-      <Flex columnGap="1rem" alignItems="center">
+      <Flex alignItems="center">
         <Text fontSize="lg" fontWeight="bold">
           {openModel}
         </Text>
-        <Text>{openedModel?.status}</Text>
       </Flex>
-      <Flex flexDirection="column">
+      <Flex
-        <Text>{openedModel?.config}</Text>
+        flexDirection="column"
-        <Text>{openedModel?.description}</Text>
+        maxHeight={window.innerHeight - 270}
-        <Text>{openedModel?.height}</Text>
+        overflowY="scroll"
-        <Text>{openedModel?.width}</Text>
+        paddingRight="2rem"
-        <Text>{openedModel?.default}</Text>
+      >
-        <Text>{openedModel?.vae}</Text>
+        <Formik
-        <Text>{openedModel?.weights}</Text>
+          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>
   );
 }