Model Manager UI 3.0 (#3778)

This PR completely ports over the Model Manager to 3.0 -- all of the
functionality has now been restored in addition to the following
changes.

- Model Manager now has been moved to its own tab on the left hand side.
- Model Manager has three tabs - Model Manager, Import Models and Merge
Models
- The edit forms for the Models now allow the users to update the model
name and the base model too along with other details.
- Checkpoint Edit form now displays the available config files from
InvokeAI and also allows users to supply their own custom config file.
- Under Import Models you can directly add models or a scan a folder for
your checkpoint files.
- Adding models has two modes -- Simple and Advanced.
- In Simple Mode, you just simply need to pass a path and InvokeAI will
try to determine kind of model it is and fill up the rest of the details
accordingly. This input lets you supply local paths to diffusers / local
paths to checkpoints / huggingface repo ID's to download models /
CivitAI links.
- Simple Mode also allows you to download different models types like
VAE's and Controlnet models and etc. Not just main models.
- In cases where the auto detection system of InvokeAI fails to read a
model correctly, you can take the manual approach and go to Advanced
where you can configure your model while adding it exactly the way you
want it. Both Diffusers and Checkpoint models now have their own custom
forms.
- Scan Models has been cleaned up. It will now only display the models
that are not already installed to InvokeAI. And each item will have two
options - Quick Add and Advanced .. replicating the Add Model behavior
from above.
- Scan Models now has a search bar for you to search through your
scanned models.
- Merge Models functionality has been restored.

This is a wrap for this PR.

**TODO: (Probably for 3.1)** 

- Add model management for model types such as VAE's and ControlNet
Models
- Replace the VAE slot on the edit forms with the installed VAE drop
down + custom option

Author: lstein

invokeai/app/api/routers/models.py

@@ -63,20 +63,35 @@ async def update_model(
 ) -> UpdateModelResponse:
     """ Update model contents with a new config. If the model name or base fields are changed, then the model is renamed. """
     logger = ApiDependencies.invoker.services.logger
+
+    
     try:
+        previous_info = ApiDependencies.invoker.services.model_manager.list_model(
+            model_name=model_name,
+            base_model=base_model,
+            model_type=model_type,
+        )
+
         # rename operation requested
         if info.model_name != model_name or info.base_model != base_model:
-            result = ApiDependencies.invoker.services.model_manager.rename_model(
+            ApiDependencies.invoker.services.model_manager.rename_model(
                 base_model = base_model,
                 model_type = model_type,
                 model_name = model_name,
                 new_name = info.model_name,
                 new_base = info.base_model,
             )
-            logger.debug(f'renaming result = {result}')
             logger.info(f'Successfully renamed {base_model}/{model_name}=>{info.base_model}/{info.model_name}')
+            # update information to support an update of attributes
             model_name = info.model_name
             base_model = info.base_model
+            new_info = ApiDependencies.invoker.services.model_manager.list_model(
+                model_name=model_name,
+                base_model=base_model,
+                model_type=model_type,
+            )
+            if new_info.get('path') != previous_info.get('path'):  # model manager moved model path during rename - don't overwrite it
+                info.path = new_info.get('path')
 
         ApiDependencies.invoker.services.model_manager.update_model(
             model_name=model_name,

invokeai/backend/model_management/model_manager.py

@@ -568,6 +568,9 @@ def list_models(
                 model_type=cur_model_type,
             )
 
+            # expose paths as absolute to help web UI
+            if path := model_dict.get('path'):
+                model_dict['path'] = str(self.app_config.root_path / path)
             models.append(model_dict)
 
         return models
@@ -635,6 +638,10 @@ def add_model(
         The returned dict has the same format as the dict returned by
         model_info().
         """
+        # relativize paths as they go in - this makes it easier to move the root directory around
+        if path := model_attributes.get('path'):
+            if Path(path).is_relative_to(self.app_config.root_path):
+                model_attributes['path'] = str(Path(path).relative_to(self.app_config.root_path))
 
         model_class = MODEL_CLASSES[base_model][model_type]
         model_config = model_class.create_config(**model_attributes)
@@ -700,7 +707,7 @@ def rename_model(
 
         # if this is a model file/directory that we manage ourselves, we need to move it
         if old_path.is_relative_to(self.app_config.models_path):
-            new_path = self.app_config.root_path / 'models' / new_base.value / model_type.value / new_name
+            new_path = self.app_config.root_path / 'models' / BaseModelType(new_base).value / ModelType(model_type).value / new_name
             move(old_path, new_path)
             model_cfg.path = str(new_path.relative_to(self.app_config.root_path))
 

invokeai/backend/model_management/models/vae.py

@@ -16,6 +16,7 @@
     calc_model_size_by_data,
     classproperty,
     InvalidModelException,
+    ModelNotFoundException,
 )
 from invokeai.app.services.config import InvokeAIAppConfig
 from diffusers.utils import is_safetensors_available

invokeai/frontend/web/public/locales/en.json

@@ -399,6 +399,8 @@
         "deleteModel": "Delete Model",
         "deleteConfig": "Delete Config",
         "deleteMsg1": "Are you sure you want to delete this model from InvokeAI?",
+        "modelDeleted": "Model Deleted",
+        "modelDeleteFailed": "Failed to delete model",
         "deleteMsg2": "This WILL delete the model from disk if it is in the InvokeAI root folder. If you are using a custom location, then the model WILL NOT be deleted from disk.",
         "formMessageDiffusersModelLocation": "Diffusers Model Location",
         "formMessageDiffusersModelLocationDesc": "Please enter at least one.",
@@ -408,11 +410,13 @@
         "convertToDiffusers": "Convert To Diffusers",
         "convertToDiffusersHelpText1": "This model will be converted to the 🧨 Diffusers format.",
         "convertToDiffusersHelpText2": "This process will replace your Model Manager entry with the Diffusers version of the same model.",
-        "convertToDiffusersHelpText3": "Your checkpoint file on the disk will NOT be deleted or modified in anyway. You can add your checkpoint to the Model Manager again if you want to.",
+        "convertToDiffusersHelpText3": "Your checkpoint file on disk WILL be deleted if it is in InvokeAI root folder. If it is in a custom location, then it WILL NOT be deleted.",
         "convertToDiffusersHelpText4": "This is a one time process only. It might take around 30s-60s depending on the specifications of your computer.",
         "convertToDiffusersHelpText5": "Please make sure you have enough disk space. Models generally vary between 2GB-7GB in size.",
         "convertToDiffusersHelpText6": "Do you wish to convert this model?",
         "convertToDiffusersSaveLocation": "Save Location",
+        "noCustomLocationProvided": "No Custom Location Provided",
+        "convertingModelBegin": "Converting Model. Please wait.",
         "v1": "v1",
         "v2_base": "v2 (512px)",
         "v2_768": "v2 (768px)",
@@ -450,7 +454,8 @@
         "none": "none",
         "addDifference": "Add Difference",
         "pickModelType": "Pick Model Type",
-        "selectModel": "Select Model"
+        "selectModel": "Select Model",
+        "importModels": "Import Models"
     },
     "parameters": {
         "general": "General",

invokeai/frontend/web/src/app/store/store.ts

@@ -21,6 +21,7 @@ import generationReducer from 'features/parameters/store/generationSlice';
 import postprocessingReducer from 'features/parameters/store/postprocessingSlice';
 import configReducer from 'features/system/store/configSlice';
 import systemReducer from 'features/system/store/systemSlice';
+import modelmanagerReducer from 'features/ui/components/tabs/ModelManager/store/modelManagerSlice';
 import hotkeysReducer from 'features/ui/store/hotkeysSlice';
 import uiReducer from 'features/ui/store/uiSlice';
 
@@ -49,6 +50,7 @@ const allReducers = {
   dynamicPrompts: dynamicPromptsReducer,
   imageDeletion: imageDeletionReducer,
   lora: loraReducer,
+  modelmanager: modelmanagerReducer,
   [api.reducerPath]: api.reducer,
 };
 
@@ -67,6 +69,7 @@ const rememberedKeys: (keyof typeof allReducers)[] = [
   'controlNet',
   'dynamicPrompts',
   'lora',
+  'modelmanager',
 ];
 
 export const store = configureStore({

invokeai/frontend/web/src/common/components/IAIInput.tsx

@@ -8,19 +8,34 @@ import {
 import { useAppDispatch } from 'app/store/storeHooks';
 import { stopPastePropagation } from 'common/util/stopPastePropagation';
 import { shiftKeyPressed } from 'features/ui/store/hotkeysSlice';
-import { ChangeEvent, KeyboardEvent, memo, useCallback } from 'react';
+import {
+  CSSProperties,
+  ChangeEvent,
+  KeyboardEvent,
+  memo,
+  useCallback,
+} from 'react';
 
 interface IAIInputProps extends InputProps {
   label?: string;
+  labelPos?: 'top' | 'side';
   value?: string;
   size?: string;
   onChange?: (e: ChangeEvent<HTMLInputElement>) => void;
   formControlProps?: Omit<FormControlProps, 'isInvalid' | 'isDisabled'>;
 }
 
+const labelPosVerticalStyle: CSSProperties = {
+  display: 'flex',
+  flexDirection: 'row',
+  alignItems: 'center',
+  gap: 10,
+};
+
 const IAIInput = (props: IAIInputProps) => {
   const {
     label = '',
+    labelPos = 'top',
     isDisabled = false,
     isInvalid,
     formControlProps,
@@ -51,6 +66,7 @@ const IAIInput = (props: IAIInputProps) => {
       isInvalid={isInvalid}
       isDisabled={isDisabled}
       {...formControlProps}
+      style={labelPos === 'side' ? labelPosVerticalStyle : undefined}
     >
       {label !== '' && <FormLabel>{label}</FormLabel>}
       <Input

invokeai/frontend/web/src/common/components/IAIMantineInput.tsx

@@ -36,6 +36,7 @@ export default function IAIMantineTextInput(props: IAIMantineTextInputProps) {
         label: {
           color: mode(base700, base300)(colorMode),
           fontWeight: 'normal',
+          marginBottom: 4,
         },
       })}
       {...rest}

invokeai/frontend/web/src/common/components/IAIMantineSelect.tsx

@@ -9,14 +9,14 @@ export type IAISelectDataType = {
   tooltip?: string;
 };
 
-type IAISelectProps = Omit<SelectProps, 'label'> & {
+export type IAISelectProps = Omit<SelectProps, 'label'> & {
   tooltip?: string;
   inputRef?: RefObject<HTMLInputElement>;
   label?: string;
 };
 
 const IAIMantineSelect = (props: IAISelectProps) => {
-  const { tooltip, inputRef, label, disabled, ...rest } = props;
+  const { tooltip, inputRef, label, disabled, required, ...rest } = props;
 
   const styles = useMantineSelectStyles();
 
@@ -25,7 +25,7 @@ const IAIMantineSelect = (props: IAISelectProps) => {
       <Select
         label={
           label ? (
-            <FormControl isDisabled={disabled}>
+            <FormControl isRequired={required} isDisabled={disabled}>
               <FormLabel>{label}</FormLabel>
             </FormControl>
           ) : undefined

invokeai/frontend/web/src/features/gallery/components/ImageGrid/GalleryImageGrid.tsx

@@ -16,14 +16,13 @@ import {
   ASSETS_CATEGORIES,
   IMAGE_CATEGORIES,
   IMAGE_LIMIT,
-  selectImagesAll,
 } from 'features/gallery//store/gallerySlice';
 import { selectFilteredImages } from 'features/gallery/store/gallerySelectors';
 import { VirtuosoGrid } from 'react-virtuoso';
 import { receivedPageOfImages } from 'services/api/thunks/image';
+import { useListBoardImagesQuery } from '../../../../services/api/endpoints/boardImages';
 import ImageGridItemContainer from './ImageGridItemContainer';
 import ImageGridListContainer from './ImageGridListContainer';
-import { useListBoardImagesQuery } from '../../../../services/api/endpoints/boardImages';
 
 const selector = createSelector(
   [stateSelector, selectFilteredImages],
@@ -180,7 +179,6 @@ const GalleryImageGrid = () => {
       </Box>
     );
   }
-  console.log({ selectedBoardId });
 
   if (status !== 'rejected') {
     return (

invokeai/frontend/web/src/features/system/store/systemSlice.ts

@@ -1,11 +1,10 @@
 import { UseToastOptions } from '@chakra-ui/react';
 import { PayloadAction, createSlice } from '@reduxjs/toolkit';
-import * as InvokeAI from 'app/types/invokeai';
 
 import { InvokeLogLevel } from 'app/logging/useLogger';
 import { userInvoked } from 'app/store/actions';
 import { nodeTemplatesBuilt } from 'features/nodes/store/nodesSlice';
-import { TFuncKey, t } from 'i18next';
+import { t } from 'i18next';
 import { LogLevelName } from 'roarr';
 import { imageUploaded } from 'services/api/thunks/image';
 import {
@@ -44,8 +43,6 @@ export interface SystemState {
   isCancelable: boolean;
   enableImageDebugging: boolean;
   toastQueue: UseToastOptions[];
-  searchFolder: string | null;
-  foundModels: InvokeAI.FoundModel[] | null;
   /**
    * The current progress image
    */
@@ -79,7 +76,7 @@ export interface SystemState {
    */
   consoleLogLevel: InvokeLogLevel;
   shouldLogToConsole: boolean;
-  statusTranslationKey: TFuncKey;
+  statusTranslationKey: any;
   /**
    * When a session is canceled, its ID is stored here until a new session is created.
    */
@@ -106,8 +103,6 @@ export const initialSystemState: SystemState = {
   isCancelable: true,
   enableImageDebugging: false,
   toastQueue: [],
-  searchFolder: null,
-  foundModels: null,
   progressImage: null,
   shouldAntialiasProgressImage: false,
   sessionId: null,
@@ -132,7 +127,7 @@ export const systemSlice = createSlice({
     setIsProcessing: (state, action: PayloadAction<boolean>) => {
       state.isProcessing = action.payload;
     },
-    setCurrentStatus: (state, action: PayloadAction<TFuncKey>) => {
+    setCurrentStatus: (state, action: any) => {
       state.statusTranslationKey = action.payload;
     },
     setShouldConfirmOnDelete: (state, action: PayloadAction<boolean>) => {
@@ -153,15 +148,6 @@ export const systemSlice = createSlice({
     clearToastQueue: (state) => {
       state.toastQueue = [];
     },
-    setSearchFolder: (state, action: PayloadAction<string | null>) => {
-      state.searchFolder = action.payload;
-    },
-    setFoundModels: (
-      state,
-      action: PayloadAction<InvokeAI.FoundModel[] | null>
-    ) => {
-      state.foundModels = action.payload;
-    },
     /**
      * A cancel was scheduled
      */
@@ -426,8 +412,6 @@ export const {
   setEnableImageDebugging,
   addToast,
   clearToastQueue,
-  setSearchFolder,
-  setFoundModels,
   cancelScheduled,
   scheduledCancelAborted,
   cancelTypeChanged,