add interactive configuration to the model loader

- Loader is renamed `configure_invokeai.py`, but `preload_models.py` is retained
  (as a shell) for backward compatibility

- At startup, if no runtime root directory exists and no `.invokeai` startup file is
  present, user will be prompted to select the runtime and outputs directories.

- Also expanded the number of initial models offered to the user to include the
  most "liked" ones from HuggingFace, including the two trinart models, the
  PaperCut model, and the VoxelArt model.

- Created a configuration file for initial models to be offered to the user, at
  configs/INITIAL_MODELS.yaml

configs/INITIAL_MODELS.yaml

@@ -0,0 +1,80 @@
+stable-diffusion-1.5:
+   description: The newest Stable Diffusion version 1.5 weight file (4.27 GB)
+   repo_id: runwayml/stable-diffusion-v1-5
+   config: v1-inference.yaml
+   file: v1-5-pruned-emaonly.ckpt
+   recommended: true
+   width: 512
+   height: 512
+inpainting-1.5:
+   description: RunwayML SD 1.5 model optimized for inpainting (4.27 GB)
+   repo_id: runwayml/stable-diffusion-inpainting
+   config: v1-inpainting-inference.yaml
+   file: sd-v1-5-inpainting.ckpt
+   recommended: True
+   width: 512
+   height: 512
+ft-mse-improved-autoencoder-840000:
+   description: StabilityAI improved autoencoder fine-tuned for human faces (recommended; 335 MB)
+   repo_id: stabilityai/sd-vae-ft-mse-original
+   config: VAE/default
+   file: vae-ft-mse-840000-ema-pruned.ckpt
+   recommended: True
+   width: 512
+   height: 512
+stable-diffusion-1.4:
+   description: The original Stable Diffusion version 1.4 weight file (4.27 GB)
+   repo_id: CompVis/stable-diffusion-v-1-4-original
+   config: v1-inference.yaml
+   file: sd-v1-4.ckpt
+   recommended: False
+   width: 512
+   height: 512
+waifu-diffusion-1.3:
+   description: Stable Diffusion 1.4 fine tuned on anime-styled images (4.27)
+   repo_id: hakurei/waifu-diffusion-v1-3
+   config: v1-inference.yaml
+   file: model-epoch09-float32.ckpt
+   recommended: False
+   width: 512
+   height: 512
+trinart-2.0:
+   description: An SD model finetuned with ~40,000 assorted high resolution manga/anime-style pictures (2.13 GB)
+   repo_id: naclbit/trinart_stable_diffusion_v2
+   config: v1-inference.yaml
+   file: trinart2_step95000.ckpt
+   recommended: False
+   width: 512
+   height: 512
+trinart_characters-1.0:
+   description: An SD model finetuned with 19.2M anime/manga style images (2.13 GB)
+   repo_id: naclbit/trinart_characters_19.2m_stable_diffusion_v1
+   config: v1-inference.yaml
+   file: trinart_characters_it4_v1.ckpt
+   recommended: False
+   width: 512
+   height: 512
+trinart_vae:
+   description: Custom autoencoder for trinart_characters
+   repo_id: naclbit/trinart_characters_19.2m_stable_diffusion_v1
+   config: VAE/trinart
+   file: autoencoder_fix_kl-f8-trinart_characters.ckpt
+   recommended: False
+   width: 512
+   height: 512
+papercut-1.0:
+   description: SD 1.5 fine-tuned for papercut art (use "PaperCut" in your prompts) (2.13 GB)
+   repo_id: Fictiverse/Stable_Diffusion_PaperCut_Model
+   config: v1-inference.yaml
+   file: PaperCut_v1.ckpt
+   recommended: False
+   width: 512
+   height: 512
+voxel_art-1.0:
+   description: Stable Diffusion trained on voxel art (use "VoxelArt" in your prompts) (4.27 GB)
+   repo_id: Fictiverse/Stable_Diffusion_VoxelArt_Model
+   config: v1-inference.yaml
+   file: VoxelArt_v1.ckpt
+   recommended: False
+   width: 512
+   height: 512

installer/install.bat

@@ -133,7 +133,7 @@ if %errorlevel% neq 0 goto err_exit
 echo ***** Installed Python dependencies *****
 
 @rem preload the models
-call .venv\Scripts\python scripts\preload_models.py
+call .venv\Scripts\python scripts\configure_invokeai.py
 set err_msg=----- model download clone failed -----
 if %errorlevel% neq 0 goto err_exit
 

installer/install.sh

@@ -200,7 +200,7 @@ _err_exit $? _err_msg
 echo -e "\n***** Installed Python dependencies *****\n"
 
 # preload the models
-.venv/bin/python3 scripts/preload_models.py
+.venv/bin/python3 scripts/configure_invokeai.py
 _err_msg="\n----- model download clone failed -----\n"
 _err_exit $? _err_msg
 

ldm/invoke/readline.py

@@ -188,6 +188,9 @@ class Completer(object):
     def set_default_dir(self, path):
         self.default_dir=path
 
+    def set_options(self,options):
+        self.options = options
+
     def get_line(self,index):
         try:
             line = self.get_history_item(index)
@@ -283,6 +286,7 @@ class Completer(object):
             partial_path = text
         else:
             switch,partial_path  = match.groups()
+
         partial_path = partial_path.lstrip()
 
         matches = list()
@@ -308,7 +312,7 @@ class Completer(object):
             if not (node.endswith(extensions) or os.path.isdir(full_path)):
                 continue
 
-            if not full_path.startswith(path):
+            if path and not full_path.startswith(path):
                 continue
 
             if switch is None:
@@ -348,6 +352,21 @@ class DummyCompleter(Completer):
     def set_line(self,line):
         print(f'# {line}')
 
+def generic_completer(commands:list)->Completer:
+    if readline_available:
+        completer = Completer(commands,[])
+        readline.set_completer(completer.complete)
+        readline.set_pre_input_hook(completer._pre_input_hook)
+        readline.set_completer_delims(' ')
+        readline.parse_and_bind('tab: complete')
+        readline.parse_and_bind('set print-completions-horizontally off')
+        readline.parse_and_bind('set page-completions on')
+        readline.parse_and_bind('set skip-completed-text on')
+        readline.parse_and_bind('set show-all-if-ambiguous on')
+    else:
+        completer = DummyCompleter(commands)
+    return completer
+
 def get_completer(opt:Args, models=[])->Completer:
     if readline_available:
         completer = Completer(COMMANDS,models)

scripts/configure_invokeai.py

@@ -21,6 +21,7 @@ from pathlib import Path
 from getpass_asterisk import getpass_asterisk
 from transformers import CLIPTokenizer, CLIPTextModel
 from ldm.invoke.globals import Globals
+from ldm.invoke.readline import generic_completer
 
 import traceback
 import requests
@@ -34,55 +35,13 @@ transformers.logging.set_verbosity_error()
 #--------------------------globals-----------------------
 Model_dir = 'models'
 Weights_dir = 'ldm/stable-diffusion-v1/'
+Dataset_path = './configs/INITIAL_MODELS.yaml'
 Default_config_file = './configs/models.yaml'
 SD_Configs = './configs/stable-diffusion'
-Datasets = {
-    'stable-diffusion-1.5':  {
-        'description': 'The newest Stable Diffusion version 1.5 weight file (4.27 GB)',
-        'repo_id': 'runwayml/stable-diffusion-v1-5',
-        'config': 'v1-inference.yaml',
-        'file': 'v1-5-pruned-emaonly.ckpt',
-        'recommended': True,
-        'width': 512,
-        'height': 512,
-    },
-    'inpainting-1.5': {
-        'description': 'RunwayML SD 1.5 model optimized for inpainting (4.27 GB)',
-        'repo_id': 'runwayml/stable-diffusion-inpainting',
-        'config': 'v1-inpainting-inference.yaml',
-        'file': 'sd-v1-5-inpainting.ckpt',
-        'recommended': True,
-        'width': 512,
-        'height': 512,
-    },
-    'stable-diffusion-1.4': {
-        'description': 'The original Stable Diffusion version 1.4 weight file (4.27 GB)',
-        'repo_id': 'CompVis/stable-diffusion-v-1-4-original',
-        'config': 'v1-inference.yaml',
-        'file': 'sd-v1-4.ckpt',
-        'recommended': False,
-        'width': 512,
-        'height': 512,
-    },
-    'waifu-diffusion-1.3': {
-        'description': 'Stable Diffusion 1.4 fine tuned on anime-styled images (4.27)',
-        'repo_id': 'hakurei/waifu-diffusion-v1-3',
-        'config': 'v1-inference.yaml',
-        'file': 'model-epoch09-float32.ckpt',
-        'recommended': False,
-        'width': 512,
-        'height': 512,
-    },
-    'ft-mse-improved-autoencoder-840000': {
-        'description': 'StabilityAI improved autoencoder fine-tuned for human faces (recommended; 335 MB)',
-        'repo_id': 'stabilityai/sd-vae-ft-mse-original',
-        'config': 'VAE',
-        'file': 'vae-ft-mse-840000-ema-pruned.ckpt',
-        'recommended': True,
-        'width': 512,
-        'height': 512,
-    },
-}
+
+Datasets = OmegaConf.load(Dataset_path)
+completer = generic_completer(['yes','no'])
+
 Config_preamble = '''# This file describes the alternative machine learning models
 # available to InvokeAI script.
 #
@@ -120,6 +79,8 @@ Have fun!
 
 #---------------------------------------------
 def yes_or_no(prompt:str, default_yes=True):
+    completer.set_options(['yes','no'])
+    completer.complete_extensions(None)  # turn off path-completion mode
     default = "y" if default_yes else 'n'
     response = input(f'{prompt} [{default}] ') or default
     if default_yes:
@@ -141,6 +102,8 @@ You may download the recommended models (about 10GB total), select a customized
 completely skip this step.
 '''
     )
+    completer.set_options(['recommended','customized','skip'])
+    completer.complete_extensions(None)  # turn off path-completion mode
     selection = None
     while selection is None:
         choice = input('Download <r>ecommended models, <c>ustomize the list, or <s>kip this step? [r]: ')
@@ -399,15 +362,18 @@ def new_config_file_contents(successfully_downloaded:dict, config_file:str)->str
         conf = OmegaConf.create()
 
     # find the VAE file, if there is one
-    vae = None
+    vaes = {}
     default_selected = False
     
     for model in successfully_downloaded:
-        if Datasets[model]['config'] == 'VAE':
-            vae = Datasets[model]['file']
+        a = Datasets[model]['config'].split('/')
+        if a[0] != 'VAE':
+            continue
+        vae_target = a[1] if len(a)>1 else 'default'
+        vaes[vae_target] = Datasets[model]['file']
     
     for model in successfully_downloaded:
-        if Datasets[model]['config'] == 'VAE': # skip VAE entries
+        if Datasets[model]['config'].startswith('VAE'): # skip VAE entries
             continue
         stanza = conf[model] if model in conf else { }
         
@@ -417,8 +383,12 @@ def new_config_file_contents(successfully_downloaded:dict, config_file:str)->str
         stanza['width'] = Datasets[model]['width']
         stanza['height'] = Datasets[model]['height']
         stanza.pop('default',None)  # this will be set later
-        if vae:
-            stanza['vae'] = os.path.normpath(os.path.join(Model_dir,Weights_dir,vae))
+        if vaes:
+            for target in vaes:
+                if re.search(target, model, flags=re.IGNORECASE):
+                    stanza['vae'] = os.path.normpath(os.path.join(Model_dir,Weights_dir,vaes[target]))
+                else:
+                    stanza['vae'] = os.path.normpath(os.path.join(Model_dir,Weights_dir,vaes['default']))
         # BUG - the first stanza is always the default. User should select.
         if not default_selected:
             stanza['default'] = True
@@ -575,7 +545,7 @@ def get_root(root:str=None)->str:
     else:
         init_file = os.path.expanduser(Globals.initfile)
         if not os.path.exists(init_file):
-            return '.'
+            return None
 
     # if we get here, then we read from initfile
     root = None
@@ -587,31 +557,63 @@ def get_root(root:str=None)->str:
             match = re.search('--root\s*=?\s*"?([^"]+)"?',l)
             if match:
                 root = match.groups()[0]
-    root = root or '.'
-    return root.strip()
+                root = root.strip()
+    return root
 
 #-------------------------------------
-def initialize_rootdir(root:str):
+def select_root(yes_to_all:bool=False):
+    default = os.path.expanduser('~/invokeai')
+    if (yes_to_all):
+        return default
+    completer.set_default_dir(default)
+    completer.complete_extensions(())
+    completer.set_line(default)
+    return input(f"Select a directory in which to install InvokeAI's models and configuration files [{default}]: ")
+
+#-------------------------------------
+def select_outputs(root:str,yes_to_all:bool=False):
+    default = os.path.normpath(os.path.join(root,'outputs'))
+    if (yes_to_all):
+        return default
+    completer.set_default_dir(os.path.expanduser('~'))
+    completer.complete_extensions(())
+    completer.set_line(default)
+    return input('Select the default directory for image outputs [{default}]: ')
+
+#-------------------------------------
+def initialize_rootdir(root:str,yes_to_all:bool=False):
     assert os.path.exists('./configs'),'Run this script from within the top level of the InvokeAI source code directory, "InvokeAI"'
-    print(f'** INITIALIZING INVOKEAI ROOT DIRECTORY **')
-    print(f'Creating a directory named {root} to contain InvokeAI models, configuration files and outputs.')
-    print(f'If you move this directory, please change its location using the --root option in "{Globals.initfile},')
-    print(f'or set the environment variable INVOKEAI_ROOT to the new location.\n')
-    for name in ('models','configs','outputs','scripts','frontend/dist'):
+    
+    print(f'** INITIALIZING INVOKEAI RUNTIME DIRECTORY **')
+    root = root or select_root(yes_to_all)
+    outputs = select_outputs(root,yes_to_all)
+    Globals.root = root
+        
+    print(f'InvokeAI models and configuration files will be placed into {root} and image outputs will be placed into {outputs}.')
+    print(f'\nYou may change these values at any time by editing the --root and --output_dir options in "{Globals.initfile}",')
+    print(f'You may also change the runtime directory by setting the environment variable INVOKEAI_ROOT.\n')
+    for name in ('models','configs','scripts','frontend/dist'):
         os.makedirs(os.path.join(root,name), exist_ok=True)
     for src in ('configs','scripts','frontend/dist'):
         dest = os.path.join(root,src)
         if not os.path.samefile(src,dest):
             shutil.copytree(src,dest,dirs_exist_ok=True)
+    os.makedirs(outputs)
 
     init_file = os.path.expanduser(Globals.initfile)
     if not os.path.exists(init_file):
-        print(f'Creating a basic initialization file at "{init_file}".\n')
+        print(f'Creating the initialization file at "{init_file}".\n')
         with open(init_file,'w') as f:
             f.write(f'''# InvokeAI initialization file
+# This is the InvokeAI initialization file, which contains command-line default values.
+# Feel free to edit. If anything goes wrong, you can re-initialize this file by deleting
+# or renaming it and then running configure_invokeai.py again.
+
 # The --root option below points to the folder in which InvokeAI stores its models, configs and outputs.
-# Don't change it unless you know what you are doing!
---root="{Globals.root}"
+--root="{root}"
+
+# the --outdir option controls the default location of image files.
+--outdir="{outputs}"
 
 # You may place other  frequently-used startup commands here, one or more per line.
 # Examples:
@@ -622,6 +624,7 @@ def initialize_rootdir(root:str):
 '''
             )
     
+    
 #-------------------------------------
 class ProgressBar():
     def __init__(self,model_name='file'):
@@ -663,8 +666,9 @@ def main():
                         help='path to root of install directory')
     opt = parser.parse_args()
 
+
     # setting a global here
-    Globals.root = os.path.expanduser(get_root(opt.root))
+    Globals.root = os.path.expanduser(get_root(opt.root) or '')
 
     try:
         introduction()
@@ -672,9 +676,12 @@ def main():
         # We check for two files to see if the runtime directory is correctly initialized.
         # 1. a key stable diffusion config file
         # 2. the web front end static files
-        if not os.path.exists(os.path.join(Globals.root,'configs/stable-diffusion/v1-inference.yaml')) \
+        if Globals.root == '' \
+           or not os.path.exists(os.path.join(Globals.root,'configs/stable-diffusion/v1-inference.yaml')) \
            or not os.path.exists(os.path.join(Globals.root,'frontend/dist')):
-            initialize_rootdir(Globals.root)
+            initialize_rootdir(Globals.root,opt.yes_to_all)
+
+        print(f'(Initializing with runtime root {Globals.root})\n')
 
         if opt.interactive:
             print('** DOWNLOADING DIFFUSION WEIGHTS **')

scripts/preload_models.py

@@ -5,8 +5,8 @@
 # two machines must share a common .cache directory.
 
 import warnings
-import load_models
+import configure_invokeai
 
 if __name__ == '__main__':
-    load_models.main()
+    configure_invokeai.main()