From 6523fd07ab2e713260595a74516b390c5552dcdc Mon Sep 17 00:00:00 2001
From: blessedcoolant <54517381+blessedcoolant@users.noreply.github.com>
Date: Mon, 26 Dec 2022 00:13:54 +1300
Subject: [PATCH] Model Edit Initial Implementation

---
 frontend/src/app/invokeai.d.ts                |  21 -
 .../components/ModelManager/AddModel.tsx      |   4 +-
 .../components/ModelManager/ModelEdit.tsx     | 375 +++---------------
 .../components/ModelManager/ModelList.tsx     |   2 +-
 ldm/invoke/model_cache.py                     |  20 +-
 5 files changed, 68 insertions(+), 354 deletions(-)

diff --git a/frontend/src/app/invokeai.d.ts b/frontend/src/app/invokeai.d.ts
index 4fa10962e4..af292058c2 100644
--- a/frontend/src/app/invokeai.d.ts
+++ b/frontend/src/app/invokeai.d.ts
@@ -264,24 +264,3 @@ export declare type UploadOutpaintingMergeImagePayload = {
   dataURL: string;
   name: string;
 };
-
-export declare type InvokeModelConfigProps = {
-  name: string;
-  description?: string;
-  config: string;
-  weights: string;
-  vae?: string;
-  width: number;
-  height: number;
-  default?: boolean;
-};
-
-export declare type OnFoundModelResponse = {
-  search_folder: string;
-  found_models: FoundModel[];
-};
-
-export declare type FoundModel = {
-  name: string;
-  location: string;
-};
diff --git a/frontend/src/features/system/components/ModelManager/AddModel.tsx b/frontend/src/features/system/components/ModelManager/AddModel.tsx
index 42ec59cd93..82f2e03d46 100644
--- a/frontend/src/features/system/components/ModelManager/AddModel.tsx
+++ b/frontend/src/features/system/components/ModelManager/AddModel.tsx
@@ -26,9 +26,9 @@ import {
 import React from 'react';
 import { FaPlus } from 'react-icons/fa';
 import { Field, FieldInputProps, Formik, FormikProps } from 'formik';
-import { RootState } from 'app/store';
+import type { RootState } from 'app/store';
 import { addNewModel } from 'app/socketio/actions';
-import { InvokeModelConfigProps } from 'app/invokeai';
+import type { InvokeModelConfigProps } from 'app/invokeai';
 import IAICheckbox from 'common/components/IAICheckbox';
 import IAIButton from 'common/components/IAIButton';
 import SearchModels from './SearchModels';
diff --git a/frontend/src/features/system/components/ModelManager/ModelEdit.tsx b/frontend/src/features/system/components/ModelManager/ModelEdit.tsx
index 1b70745e18..4738b0a2a7 100644
--- a/frontend/src/features/system/components/ModelManager/ModelEdit.tsx
+++ b/frontend/src/features/system/components/ModelManager/ModelEdit.tsx
@@ -1,331 +1,60 @@
-import { createSelector } from '@reduxjs/toolkit';
-import { Field, FieldInputProps, Formik, FormikProps } from 'formik';
-import { RootState } from 'app/store';
-import { useAppDispatch, useAppSelector } from 'app/storeHooks'
-import { addNewModel,  } from 'app/socketio/actions';
-import { InvokeModelConfigProps } from 'app/invokeai';
-import { SystemState } from 'features/system/store/systemSlice';
-import {
-  Button,
-  Flex,
-  FormControl,
-  FormErrorMessage,
-  FormHelperText,
-  FormLabel,
-  HStack,
-  Input,
-  NumberDecrementStepper,
-  NumberIncrementStepper,
-  NumberInput,
-  NumberInputField,
-  NumberInputStepper,
-  Text,
-  VStack,
-} from '@chakra-ui/react';
-
-
+import { createSelector, Dictionary } from '@reduxjs/toolkit';
+import type { RootState } from 'app/store';
+import { 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';
 
+const selector = createSelector(
+  [systemSelector],
+  (system) => {
+    const { openModel, model_list } = system;
+    return {
+      model_list,
+      openModel,
+    };
+  },
+  {
+    memoizeOptions: {
+      resultEqualityCheck: _.isEqual,
+    },
+  }
+);
 
 export default function ModelEdit() {
+  const { openModel, model_list } = useAppSelector(selector);
+  const [openedModel, setOpenedModel] = useState<Model>();
 
-  const MIN_MODEL_SIZE = 64;
-  const MAX_MODEL_SIZE = 2048;
-  const modelListSelector = createSelector(
-    (state: RootState) => state.system,
-    (system: SystemState) => {
-      const models = _.map(system.model_list, (model, key) => {
-        return { name: key, ...model };
+  useEffect(() => {
+    if (openModel) {
+      const retrievedModel = _.pickBy(model_list, (val, key) => {
+        return _.isEqual(key, openModel);
       });
-  
-      const selectedModel = models.find((model) => model.name === system.openModel);
-     
-      return {
-        openedModel: {
-          name:selectedModel?.name,
-          description:selectedModel?.description,
-          config:selectedModel?.config,
-          weights:selectedModel?.weights,
-          vae:selectedModel?.vae,
-          width:selectedModel?.width,
-          height:selectedModel?.height,
-          default:selectedModel?.default,
-        }
-      };
+      setOpenedModel(retrievedModel[openModel]);
     }
+  }, [model_list, openModel]);
+
+  console.log(openedModel);
+
+  return (
+    <Flex flexDirection="column" rowGap="1rem">
+      <Flex columnGap="1rem" 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>
+    </Flex>
   );
-  
- 
-  const { openedModel } = useAppSelector(modelListSelector);
-
-
-  const viewModelFormValues: InvokeModelConfigProps = {
-    name: openedModel.name || '',
-    description: openedModel.description || '',
-    config: openedModel.config || '',
-    weights: openedModel.weights || '',
-    vae: openedModel.vae || '',
-    width: openedModel.width || 0,
-    height: openedModel.height || 0,
-    default: openedModel.default || false,
-  };
-
-  const dispatch = useAppDispatch();
-  
-  const addModelFormSubmitHandler = (values: InvokeModelConfigProps) => {
-    dispatch(addNewModel(values));
-    onClose();
-  };
-
-  function hasWhiteSpace(s: string) {
-    return /\\s/g.test(s);
-  }
-
-  function baseValidation(value: string) {
-    let error;
-    if (hasWhiteSpace(value)) error = 'Cannot use spaces';
-    return error;
-  }
-
-  if (!openedModel) {
-    return ( 
-      <Flex> 'No model selected' </Flex>
-    
-  )} else {
-    
-    return (
-    <Formik
-      initialValues={viewModelFormValues}
-      onSubmit={addModelFormSubmitHandler}
-    >
-      {({ handleSubmit, errors, touched }) => (
-        <form onSubmit={handleSubmit}>
-          <VStack rowGap={'0.5rem'}>
-            <Text fontSize={20} fontWeight="bold" alignSelf={'start'}>
-              Manual
-            </Text>
-            {/* Name */}
-            <FormControl
-              isInvalid={!!errors.name && touched.name}
-              isRequired
-            >
-              <FormLabel htmlFor="name">Name</FormLabel>
-              <VStack alignItems={'start'}>
-                <Field
-                  as={Input}
-                  id="name"
-                  name="name"
-                  type="text"
-                  validate={baseValidation}
-                  readOnly
-                />
-                {!!errors.name && touched.name ? (
-                  <FormErrorMessage>{errors.name}</FormErrorMessage>
-                ) : (
-                  <FormHelperText margin={0}>
-                    Enter a name for your model
-                  </FormHelperText>
-                )}
-              </VStack>
-            </FormControl>
-
-            {/* Description */}
-            <FormControl
-              isInvalid={!!errors.description && touched.description}
-              isRequired
-            >
-              <FormLabel htmlFor="description">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}>
-                    Add a description for your model
-                  </FormHelperText>
-                )}
-              </VStack>
-            </FormControl>
-
-            {/* Config */}
-            <FormControl
-              isInvalid={!!errors.config && touched.config}
-              isRequired
-            >
-              <FormLabel htmlFor="config">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}>
-                    Path to the config file of your model.
-                  </FormHelperText>
-                )}
-              </VStack>
-            </FormControl>
-
-            {/* Weights */}
-            <FormControl
-              isInvalid={!!errors.weights && touched.weights}
-              isRequired
-            >
-              <FormLabel htmlFor="config">Model Location</FormLabel>
-              <VStack alignItems={'start'}>
-                <Field
-                  as={Input}
-                  id="weights"
-                  name="weights"
-                  type="text"
-                />
-                {!!errors.weights && touched.weights ? (
-                  <FormErrorMessage>
-                    {errors.weights}
-                  </FormErrorMessage>
-                ) : (
-                  <FormHelperText margin={0}>
-                    Path to where your model is located.
-                  </FormHelperText>
-                )}
-              </VStack>
-            </FormControl>
-
-            {/* VAE */}
-            <FormControl isInvalid={!!errors.vae && touched.vae}>
-              <FormLabel htmlFor="vae">VAE Location</FormLabel>
-              <VStack alignItems={'start'}>
-                <Field as={Input} id="vae" name="vae" type="text" />
-                {!!errors.vae && touched.vae ? (
-                  <FormErrorMessage>{errors.vae}</FormErrorMessage>
-                ) : (
-                  <FormHelperText margin={0}>
-                    Path to where your VAE is located.
-                  </FormHelperText>
-                )}
-              </VStack>
-            </FormControl>
-
-            <HStack width={'100%'}>
-              {/* Width */}
-              <FormControl
-                isInvalid={!!errors.width && touched.width}
-              >
-                <FormLabel htmlFor="width">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}>
-                      Default width of your model.
-                    </FormHelperText>
-                  )}
-                </VStack>
-              </FormControl>
-
-              {/* Height */}
-              <FormControl
-                isInvalid={!!errors.height && touched.height}
-              >
-                <FormLabel htmlFor="height">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}>
-                      Default height of your model.
-                    </FormHelperText>
-                  )}
-                </VStack>
-              </FormControl>
-            </HStack>
-
-            <Button
-              type="submit"
-              className="modal-close-btn"
-              isLoading={isProcessing}
-            >
-              Save Model
-            </Button>
-          </VStack>
-        </form>
-      )}
-    </Formik>
-)}
 }
diff --git a/frontend/src/features/system/components/ModelManager/ModelList.tsx b/frontend/src/features/system/components/ModelManager/ModelList.tsx
index 8569ad3e62..6e118d1a9f 100644
--- a/frontend/src/features/system/components/ModelManager/ModelList.tsx
+++ b/frontend/src/features/system/components/ModelManager/ModelList.tsx
@@ -68,7 +68,7 @@ const ModelList = () => {
   };
 
   return (
-    <Flex flexDirection={'column'} rowGap="2rem" width={'50%'}>
+    <Flex flexDirection={'column'} rowGap="2rem" width="45%" minWidth="45%">
       <Flex justifyContent={'space-between'}>
         <Text fontSize={'1.4rem'} fontWeight="bold">
           {t('modelmanager:availableModels')}
diff --git a/ldm/invoke/model_cache.py b/ldm/invoke/model_cache.py
index 2f670eac62..623160e67d 100644
--- a/ldm/invoke/model_cache.py
+++ b/ldm/invoke/model_cache.py
@@ -134,12 +134,13 @@ class ModelCache(object):
         '''
         models = {}
         for name in self.config:
-            try:
-                description = self.config[name].description
-                weights = self.config[name].weights
-            except ConfigAttributeError:
-                description = '<no description>'
-                weights = '<not found>'
+            description = self.config[name].description if 'description' in self.config[name] else '<no description>'
+            weights = self.config[name].weights if 'weights' in self.config[name] else '<no weights>'
+            config = self.config[name].config if 'config' in self.config[name] else '<no config>'
+            width = self.config[name].width if 'width' in self.config[name] else 512
+            height = self.config[name].height if 'height' in self.config[name] else 512
+            default = self.config[name].default if 'default' in self.config[name] else False
+            vae = self.config[name].vae if 'vae' in self.config[name] else '<no vae>'
 
             if self.current_model == name:
                 status = 'active'
@@ -151,7 +152,12 @@ class ModelCache(object):
             models[name]={
                 'status' : status,
                 'description' : description,
-                'weights': weights
+                'weights': weights,
+                'config': config,
+                'width': width,
+                'height': height,
+                'vae': vae,
+                'default': default
             }
         return models