Merge branch 'main' into diffusers-upgrade

invokeai/backend/install/invokeai_configure.py

@@ -12,8 +12,8 @@ print("Loading Python libraries...\n",file=sys.stderr)
 import argparse
 import io
 import os
-import re
 import shutil
+import textwrap
 import traceback
 import warnings
 from argparse import Namespace
@@ -38,7 +38,6 @@ from transformers import (
 import invokeai.configs as configs
 
 from invokeai.app.services.config import (
-    get_invokeai_config,
     InvokeAIAppConfig,
 )
 from invokeai.frontend.install.model_install import addModelsForm, process_and_execute
@@ -46,20 +45,20 @@ from invokeai.frontend.install.widgets import (
     CenteredButtonPress,
     IntTitleSlider,
     set_min_terminal_size,
+    CyclingForm,
+    MIN_COLS,
+    MIN_LINES,
 )
-
-from invokeai.backend.config.legacy_arg_parsing import legacy_parser
-from invokeai.backend.config.model_install_backend import (
+from invokeai.backend.install.legacy_arg_parsing import legacy_parser
+from invokeai.backend.install.model_install_backend import (
     default_dataset,
     download_from_hf,
     hf_download_with_resume,
     recommended_datasets,
+    UserSelections,
 )
 
-from invokeai.app.services.config import InvokeAIAppConfig
-
 warnings.filterwarnings("ignore")
-
 transformers.logging.set_verbosity_error()
 
 
@@ -70,18 +69,9 @@ config = InvokeAIAppConfig.get_config()
 Model_dir = "models"
 Weights_dir = "ldm/stable-diffusion-v1/"
 
-# the initial "configs" dir is now bundled in the `invokeai.configs` package
-Dataset_path = Path(configs.__path__[0]) / "INITIAL_MODELS.yaml"
-
 Default_config_file = config.model_conf_path
 SD_Configs = config.legacy_conf_path
 
-Datasets = OmegaConf.load(Dataset_path)
-
-# minimum size for the UI
-MIN_COLS = 135
-MIN_LINES = 45
-
 PRECISION_CHOICES = ['auto','float16','float32','autocast']
 
 INIT_FILE_PREAMBLE = """# InvokeAI initialization file
@@ -106,7 +96,7 @@ Command-line client:
    invokeai
 
 If you installed using an installation script, run:
-  {config.root}/invoke.{"bat" if sys.platform == "win32" else "sh"}
+  {config.root_path}/invoke.{"bat" if sys.platform == "win32" else "sh"}
 
 Add the '--help' argument to see all of the command-line switches available for use.
 """
@@ -218,16 +208,11 @@ def download_realesrgan():
     model_url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth"
     wdn_model_url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth"
 
-    model_dest = os.path.join(
-        config.root, "models/realesrgan/realesr-general-x4v3.pth"
-    )
+    model_dest = config.root_path / "models/realesrgan/realesr-general-x4v3.pth"
+    wdn_model_dest = config.root_path / "models/realesrgan/realesr-general-wdn-x4v3.pth"
 
-    wdn_model_dest = os.path.join(
-        config.root, "models/realesrgan/realesr-general-wdn-x4v3.pth"
-    )
-
-    download_with_progress_bar(model_url, model_dest, "RealESRGAN")
-    download_with_progress_bar(wdn_model_url, wdn_model_dest, "RealESRGANwdn")
+    download_with_progress_bar(model_url, str(model_dest), "RealESRGAN")
+    download_with_progress_bar(wdn_model_url, str(wdn_model_dest), "RealESRGANwdn")
 
 
 def download_gfpgan():
@@ -246,8 +231,8 @@ def download_gfpgan():
             "./models/gfpgan/weights/parsing_parsenet.pth",
         ],
     ):
-        model_url, model_dest = model[0], os.path.join(config.root, model[1])
-        download_with_progress_bar(model_url, model_dest, "GFPGAN weights")
+        model_url, model_dest = model[0], config.root_path / model[1]
+        download_with_progress_bar(model_url, str(model_dest), "GFPGAN weights")
 
 
 # ---------------------------------------------
@@ -256,8 +241,8 @@ def download_codeformer():
     model_url = (
         "https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth"
     )
-    model_dest = os.path.join(config.root, "models/codeformer/codeformer.pth")
-    download_with_progress_bar(model_url, model_dest, "CodeFormer")
+    model_dest = config.root_path / "models/codeformer/codeformer.pth"
+    download_with_progress_bar(model_url, str(model_dest), "CodeFormer")
 
 
 # ---------------------------------------------
@@ -309,7 +294,7 @@ def download_vaes():
         if not hf_download_with_resume(
             repo_id=repo_id,
             model_name=model_name,
-            model_dir=str(config.root / Model_dir / Weights_dir),
+            model_dir=str(config.root_path / Model_dir / Weights_dir),
         ):
             raise Exception(f"download of {model_name} failed")
     except Exception as e:
@@ -324,24 +309,24 @@ def get_root(root: str = None) -> str:
     elif os.environ.get("INVOKEAI_ROOT"):
         return os.environ.get("INVOKEAI_ROOT")
     else:
-        return config.root
+        return str(config.root_path)
 
 # -------------------------------------
-class editOptsForm(npyscreen.FormMultiPage):
+class editOptsForm(CyclingForm, npyscreen.FormMultiPage):
     # for responsive resizing - disabled
     # FIX_MINIMUM_SIZE_WHEN_CREATED = False
 
     def create(self):
         program_opts = self.parentApp.program_opts
         old_opts = self.parentApp.invokeai_opts
-        first_time = not (config.root / 'invokeai.yaml').exists()
+        first_time = not (config.root_path / 'invokeai.yaml').exists()
         access_token = HfFolder.get_token()
         window_width, window_height = get_terminal_size()
-        for i in [
-            "Configure startup settings. You can come back and change these later.",
-            "Use ctrl-N and ctrl-P to move to the <N>ext and <P>revious fields.",
-            "Use cursor arrows to make a checkbox selection, and space to toggle.",
-        ]:
+        label = """Configure startup settings. You can come back and change these later. 
+Use ctrl-N and ctrl-P to move to the <N>ext and <P>revious fields.
+Use cursor arrows to make a checkbox selection, and space to toggle.
+"""
+        for i in textwrap.wrap(label,width=window_width-6):
             self.add_widget_intelligent(
                 npyscreen.FixedText,
                 value=i,
@@ -368,7 +353,7 @@ class editOptsForm(npyscreen.FormMultiPage):
         self.outdir = self.add_widget_intelligent(
             npyscreen.TitleFilename,
             name="(<tab> autocompletes, ctrl-N advances):",
-            value=str(old_opts.outdir) or str(default_output_dir()),
+            value=str(default_output_dir()),
             select_dir=True,
             must_exist=False,
             use_two_lines=False,
@@ -391,14 +376,13 @@ class editOptsForm(npyscreen.FormMultiPage):
             scroll_exit=True,
         )
         self.nextrely += 1
-        for i in [
-            "If you have an account at HuggingFace you may optionally paste your access token here",
-            'to allow InvokeAI to download restricted styles & subjects from the "Concept Library".',
-            "See https://huggingface.co/settings/tokens",
-        ]:
+        label = """If you have an account at HuggingFace you may optionally paste your access token here
+to allow InvokeAI to download restricted styles & subjects from the "Concept Library". See https://huggingface.co/settings/tokens.
+"""
+        for line in textwrap.wrap(label,width=window_width-6):
             self.add_widget_intelligent(
                 npyscreen.FixedText,
-                value=i,
+                value=line,
                 editable=False,
                 color="CONTROL",
             )
@@ -475,7 +459,7 @@ class editOptsForm(npyscreen.FormMultiPage):
         self.nextrely += 1
         self.add_widget_intelligent(
             npyscreen.FixedText,
-            value="Directories containing textual inversion and LoRA models (<tab> autocompletes, ctrl-N advances):",
+            value="Directories containing textual inversion, controlnet and LoRA models (<tab> autocompletes, ctrl-N advances):",
             editable=False,
             color="CONTROL",
         )
@@ -501,6 +485,17 @@ class editOptsForm(npyscreen.FormMultiPage):
             begin_entry_at=32,
             scroll_exit=True,
         )
+        self.controlnet_dir = self.add_widget_intelligent(
+            npyscreen.TitleFilename,
+            name="             ControlNets:",
+            value=str(default_controlnet_dir()),
+            select_dir=True,
+            must_exist=False,
+            use_two_lines=False,
+            labelColor="GOOD",
+            begin_entry_at=32,
+            scroll_exit=True,
+        )
         self.nextrely += 1
         self.add_widget_intelligent(
             npyscreen.TitleFixedText,
@@ -511,11 +506,11 @@ class editOptsForm(npyscreen.FormMultiPage):
             scroll_exit=True,
         )
         self.nextrely -= 1
-        for i in [
-            "BY DOWNLOADING THE STABLE DIFFUSION WEIGHT FILES, YOU AGREE TO HAVE READ",
-            "AND ACCEPTED THE CREATIVEML RESPONSIBLE AI LICENSE LOCATED AT",
-            "https://huggingface.co/spaces/CompVis/stable-diffusion-license",
-        ]:
+        label = """BY DOWNLOADING THE STABLE DIFFUSION WEIGHT FILES, YOU AGREE TO HAVE READ
+AND ACCEPTED THE CREATIVEML RESPONSIBLE AI LICENSE LOCATED AT
+https://huggingface.co/spaces/CompVis/stable-diffusion-license
+"""
+        for i in textwrap.wrap(label,width=window_width-6):
             self.add_widget_intelligent(
                 npyscreen.FixedText,
                 value=i,
@@ -554,7 +549,7 @@ class editOptsForm(npyscreen.FormMultiPage):
             self.editing = False
         else:
             self.editing = True
-
+            
     def validate_field_values(self, opt: Namespace) -> bool:
         bad_fields = []
         if not opt.license_acceptance:
@@ -590,6 +585,7 @@ class editOptsForm(npyscreen.FormMultiPage):
                 "always_use_cpu",
                 "embedding_dir",
                 "lora_dir",
+                "controlnet_dir",
         ]:
             setattr(new_opts, attr, getattr(self, attr).value)
 
@@ -617,6 +613,7 @@ class EditOptApplication(npyscreen.NPSAppManaged):
             "MAIN",
             editOptsForm,
             name="InvokeAI Startup Options",
+            cycle_widgets=True,
         )
         if not (self.program_opts.skip_sd_weights or self.program_opts.default_only):
             self.model_select = self.addForm(
@@ -624,6 +621,7 @@ class EditOptApplication(npyscreen.NPSAppManaged):
                 addModelsForm,
                 name="Install Stable Diffusion Models",
                 multipage=True,
+                cycle_widgets=True,
             )
 
     def new_opts(self):
@@ -638,16 +636,13 @@ def edit_opts(program_opts: Namespace, invokeai_opts: Namespace) -> argparse.Nam
 
 def default_startup_options(init_file: Path) -> Namespace:
     opts = InvokeAIAppConfig.get_config()
-    outdir = Path(opts.outdir)
-    if not outdir.is_absolute():
-        opts.outdir = str(config.root / opts.outdir)
     if not init_file.exists():
         opts.nsfw_checker = True
     return opts
 
-def default_user_selections(program_opts: Namespace) -> Namespace:
-    return Namespace(
-        starter_models=default_dataset()
+def default_user_selections(program_opts: Namespace) -> UserSelections:
+    return UserSelections(
+        install_models=default_dataset()
         if program_opts.default_only
         else recommended_datasets()
         if program_opts.yes_to_all
@@ -655,26 +650,27 @@ def default_user_selections(program_opts: Namespace) -> Namespace:
         purge_deleted_models=False,
         scan_directory=None,
         autoscan_on_startup=None,
-        import_model_paths=None,
-        convert_to_diffusers=None,
     )
 
 
 # -------------------------------------
-def initialize_rootdir(root: str, yes_to_all: bool = False):
+def initialize_rootdir(root: Path, yes_to_all: bool = False):
     print("** INITIALIZING INVOKEAI RUNTIME DIRECTORY **")
 
     for name in (
-        "models",
-        "configs",
-        "embeddings",
-        "text-inversion-output",
-        "text-inversion-training-data",
+            "models",
+            "configs",
+            "embeddings",
+            "databases",
+            "loras",
+            "controlnets",
+            "text-inversion-output",
+            "text-inversion-training-data",
     ):
         os.makedirs(os.path.join(root, name), exist_ok=True)
 
     configs_src = Path(configs.__path__[0])
-    configs_dest = Path(root) / "configs"
+    configs_dest = root / "configs"
     if not os.path.samefile(configs_src, configs_dest):
         shutil.copytree(configs_src, configs_dest, dirs_exist_ok=True)
 
@@ -685,8 +681,17 @@ def run_console_ui(
 ) -> (Namespace, Namespace):
     # parse_args() will read from init file if present
     invokeai_opts = default_startup_options(initfile)
+    invokeai_opts.root = program_opts.root
 
-    set_min_terminal_size(MIN_COLS, MIN_LINES)
+    # The third argument is needed in the Windows 11 environment to
+    # launch a console window running this program.
+    set_min_terminal_size(MIN_COLS, MIN_LINES,'invokeai-configure')
+
+    # the install-models application spawns a subprocess to install
+    # models, and will crash unless this is set before running.
+    import torch
+    torch.multiprocessing.set_start_method("spawn")
+    
     editApp = EditOptApplication(program_opts, invokeai_opts)
     editApp.run()
     if editApp.user_cancelled:
@@ -700,27 +705,32 @@ def write_opts(opts: Namespace, init_file: Path):
     """
     Update the invokeai.yaml file with values from current settings.
     """
-
     # this will load current settings
-    config = InvokeAIAppConfig.get_config()
+    new_config = InvokeAIAppConfig.get_config()
+    new_config.root = config.root
+    
     for key,value in opts.__dict__.items():
-        if hasattr(config,key):
-            setattr(config,key,value)
+        if hasattr(new_config,key):
+            setattr(new_config,key,value)
 
     with open(init_file,'w', encoding='utf-8') as file:
-        file.write(config.to_yaml())
+        file.write(new_config.to_yaml())
 
 # -------------------------------------
 def default_output_dir() -> Path:
-    return config.root / "outputs"
+    return config.root_path / "outputs"
 
 # -------------------------------------
 def default_embedding_dir() -> Path:
-    return config.root / "embeddings"
+    return config.root_path / "embeddings"
 
 # -------------------------------------
 def default_lora_dir() -> Path:
-    return config.root / "loras"
+    return config.root_path / "loras"
+
+# -------------------------------------
+def default_controlnet_dir() -> Path:
+    return config.root_path / "controlnets"
 
 # -------------------------------------
 def write_default_options(program_opts: Namespace, initfile: Path):
@@ -808,9 +818,12 @@ def main():
     )
     opt = parser.parse_args()
 
-    # setting a global here
-    global config
-    config.root = Path(os.path.expanduser(get_root(opt.root) or ""))
+    invoke_args = []
+    if opt.root:
+        invoke_args.extend(['--root',opt.root])
+    if opt.full_precision:
+        invoke_args.extend(['--precision','float32'])
+    config.parse_args(invoke_args)
 
     errors = set()
 
@@ -818,17 +831,16 @@ def main():
         models_to_download = default_user_selections(opt)
 
         # We check for to see if the runtime directory is correctly initialized.
-        old_init_file = Path(config.root, 'invokeai.init')
-        new_init_file = Path(config.root, 'invokeai.yaml')
+        old_init_file = config.root_path / 'invokeai.init'
+        new_init_file = config.root_path / 'invokeai.yaml'
         if old_init_file.exists() and not new_init_file.exists():
             print('** Migrating invokeai.init to invokeai.yaml')
             migrate_init_file(old_init_file)
-
             # Load new init file into config
             config.parse_args(argv=[],conf=OmegaConf.load(new_init_file))
 
         if not config.model_conf_path.exists():
-            initialize_rootdir(config.root, opt.yes_to_all)
+            initialize_rootdir(config.root_path, opt.yes_to_all)
 
         if opt.yes_to_all:
             write_default_options(opt, new_init_file)
@@ -848,7 +860,7 @@ def main():
         if opt.skip_support_models:
             print("\n** SKIPPING SUPPORT MODEL DOWNLOADS PER USER REQUEST **")
         else:
-            print("\n** DOWNLOADING SUPPORT MODELS **")
+            print("\n** CHECKING/UPDATING SUPPORT MODELS **")
             download_bert()
             download_sd1_clip()
             download_sd2_clip()
@@ -866,6 +878,8 @@ def main():
             process_and_execute(opt, models_to_download)
 
         postscript(errors=errors)
+        if not opt.yes_to_all:
+            input('Press any key to continue...')
     except KeyboardInterrupt:
         print("\nGoodbye! Come back soon.")
 

invokeai/backend/install/model_install_backend.py

@@ -6,23 +6,24 @@ import re
 import shutil
 import sys
 import warnings
+from dataclasses import dataclass,field
 from pathlib import Path
 from tempfile import TemporaryFile
-from typing import List
+from typing import List, Dict, Callable
 
 import requests
 from diffusers import AutoencoderKL
-from huggingface_hub import hf_hub_url
+from huggingface_hub import hf_hub_url, HfFolder
 from omegaconf import OmegaConf
 from omegaconf.dictconfig import DictConfig
 from tqdm import tqdm
 
 import invokeai.configs as configs
 
-from invokeai.app.services.config import InvokeAIAppConfig
-from ..model_management import ModelManager
-from ..stable_diffusion import StableDiffusionGeneratorPipeline
 
+from invokeai.app.services.config import InvokeAIAppConfig
+from ..stable_diffusion import StableDiffusionGeneratorPipeline
+from ..util.logging import InvokeAILogger
 
 warnings.filterwarnings("ignore")
 
@@ -38,6 +39,9 @@ Dataset_path = Path(configs.__path__[0]) / "INITIAL_MODELS.yaml"
 # initial models omegaconf
 Datasets = None
 
+# logger
+logger = InvokeAILogger.getLogger(name='InvokeAI')
+
 Config_preamble = """
 # This file describes the alternative machine learning models
 # available to InvokeAI script.
@@ -48,11 +52,30 @@ Config_preamble = """
 # was trained on.
 """
 
+@dataclass
+class ModelInstallList:
+    '''Class for listing models to be installed/removed'''
+    install_models: List[str] = field(default_factory=list)
+    remove_models: List[str] = field(default_factory=list)
 
+@dataclass
+class UserSelections():
+    install_models: List[str]= field(default_factory=list)
+    remove_models: List[str]=field(default_factory=list)
+    purge_deleted_models: bool=field(default_factory=list)
+    install_cn_models: List[str] = field(default_factory=list)
+    remove_cn_models: List[str] = field(default_factory=list)
+    install_lora_models: List[str] = field(default_factory=list)
+    remove_lora_models: List[str] = field(default_factory=list)
+    install_ti_models: List[str] = field(default_factory=list)
+    remove_ti_models: List[str] = field(default_factory=list)
+    scan_directory: Path = None
+    autoscan_on_startup: bool=False
+    import_model_paths: str=None
+        
 def default_config_file():
     return config.model_conf_path
 
-
 def sd_configs():
     return config.legacy_conf_path
 
@@ -60,45 +83,67 @@ def initial_models():
     global Datasets
     if Datasets:
         return Datasets
-    return (Datasets := OmegaConf.load(Dataset_path))
-
+    return (Datasets := OmegaConf.load(Dataset_path)['diffusers'])
 
 def install_requested_models(
-    install_initial_models: List[str] = None,
-    remove_models: List[str] = None,
-    scan_directory: Path = None,
-    external_models: List[str] = None,
-    scan_at_startup: bool = False,
-    precision: str = "float16",
-    purge_deleted: bool = False,
-    config_file_path: Path = None,
+        diffusers: ModelInstallList = None,
+        controlnet: ModelInstallList = None,
+        lora: ModelInstallList = None,
+        ti: ModelInstallList = None,
+        cn_model_map: Dict[str,str] = None, # temporary - move to model manager
+        scan_directory: Path = None,
+        external_models: List[str] = None,
+        scan_at_startup: bool = False,
+        precision: str = "float16",
+        purge_deleted: bool = False,
+        config_file_path: Path = None,
+        model_config_file_callback:  Callable[[Path],Path] = None
 ):
     """
     Entry point for installing/deleting starter models, or installing external models.
     """
+    access_token = HfFolder.get_token()
     config_file_path = config_file_path or default_config_file()
     if not config_file_path.exists():
         open(config_file_path, "w")
 
+    # prevent circular import here
+    from ..model_management import ModelManager
     model_manager = ModelManager(OmegaConf.load(config_file_path), precision=precision)
+    if controlnet:
+        model_manager.install_controlnet_models(controlnet.install_models, access_token=access_token)
+        model_manager.delete_controlnet_models(controlnet.remove_models)
 
-    if remove_models and len(remove_models) > 0:
-        print("== DELETING UNCHECKED STARTER MODELS ==")
-        for model in remove_models:
-            print(f"{model}...")
-            model_manager.del_model(model, delete_files=purge_deleted)
-        model_manager.commit(config_file_path)
+    if lora:
+        model_manager.install_lora_models(lora.install_models, access_token=access_token)
+        model_manager.delete_lora_models(lora.remove_models)
 
-    if install_initial_models and len(install_initial_models) > 0:
-        print("== INSTALLING SELECTED STARTER MODELS ==")
-        successfully_downloaded = download_weight_datasets(
-            models=install_initial_models,
-            access_token=None,
-            precision=precision,
-        )  # FIX: for historical reasons, we don't use model manager here
-        update_config_file(successfully_downloaded, config_file_path)
-        if len(successfully_downloaded) < len(install_initial_models):
-            print("** Some of the model downloads were not successful")
+    if ti:
+        model_manager.install_ti_models(ti.install_models, access_token=access_token)
+        model_manager.delete_ti_models(ti.remove_models)
+
+    if diffusers:
+        # TODO: Replace next three paragraphs with calls into new model manager
+        if diffusers.remove_models and len(diffusers.remove_models) > 0:
+            logger.info("Processing requested deletions")
+            for model in diffusers.remove_models:
+                logger.info(f"{model}...")
+                model_manager.del_model(model, delete_files=purge_deleted)
+            model_manager.commit(config_file_path)
+
+        if diffusers.install_models and len(diffusers.install_models) > 0:
+            logger.info("Installing requested models")
+            downloaded_paths = download_weight_datasets(
+                models=diffusers.install_models,
+                access_token=None,
+                precision=precision,
+            )
+            successful = {x:v for x,v in downloaded_paths.items() if v is not None}
+            if len(successful) > 0:
+                update_config_file(successful, config_file_path)
+            if len(successful) < len(diffusers.install_models):
+                unsuccessful = [x for x in downloaded_paths if downloaded_paths[x] is None]
+                logger.warning(f"Some of the model downloads were not successful: {unsuccessful}")
 
     # due to above, we have to reload the model manager because conf file
     # was changed behind its back
@@ -109,12 +154,14 @@ def install_requested_models(
         external_models.append(str(scan_directory))
 
     if len(external_models) > 0:
-        print("== INSTALLING EXTERNAL MODELS ==")
+        logger.info("INSTALLING EXTERNAL MODELS")
         for path_url_or_repo in external_models:
             try:
+                logger.debug(f'In install_requested_models; callback = {model_config_file_callback}')
                 model_manager.heuristic_import(
                     path_url_or_repo,
                     commit_to_conf=config_file_path,
+                    config_file_callback = model_config_file_callback,
                 )
             except KeyboardInterrupt:
                 sys.exit(-1)
@@ -122,18 +169,22 @@ def install_requested_models(
                 pass
 
     if scan_at_startup and scan_directory.is_dir():
-        argument = "--autoconvert"
-        print('** The global initfile is no longer supported; rewrite to support new yaml format **')
-        initfile = Path(config.root, 'invokeai.init')
-        replacement = Path(config.root, f"invokeai.init.new")
-        directory = str(scan_directory).replace("\\", "/")
-        with open(initfile, "r") as input:
-            with open(replacement, "w") as output:
-                while line := input.readline():
-                    if not line.startswith(argument):
-                        output.writelines([line])
-                output.writelines([f"{argument} {directory}"])
-        os.replace(replacement, initfile)
+        update_autoconvert_dir(scan_directory)
+    else:
+        update_autoconvert_dir(None)
+
+def update_autoconvert_dir(autodir: Path):
+    '''
+    Update the "autoconvert_dir" option in invokeai.yaml
+    '''
+    invokeai_config_path = config.init_file_path
+    conf = OmegaConf.load(invokeai_config_path)
+    conf.InvokeAI.Paths.autoconvert_dir = str(autodir) if autodir else None
+    yaml = OmegaConf.to_yaml(conf)
+    tmpfile = invokeai_config_path.parent / "new_config.tmp"
+    with open(tmpfile, "w", encoding="utf-8") as outfile:
+        outfile.write(yaml)
+    tmpfile.replace(invokeai_config_path)
 
 
 # -------------------------------------
@@ -145,33 +196,21 @@ def yes_or_no(prompt: str, default_yes=True):
     else:
         return response[0] in ("y", "Y")
 
-
-# -------------------------------------
-def get_root(root: str = None) -> str:
-    if root:
-        return root
-    elif os.environ.get("INVOKEAI_ROOT"):
-        return os.environ.get("INVOKEAI_ROOT")
-    else:
-        return config.root
-
-
 # ---------------------------------------------
-def recommended_datasets() -> dict:
-    datasets = dict()
+def recommended_datasets() -> List['str']:
+    datasets = set()
     for ds in initial_models().keys():
         if initial_models()[ds].get("recommended", False):
-            datasets[ds] = True
-    return datasets
-
+            datasets.add(ds)
+    return list(datasets)
 
 # ---------------------------------------------
 def default_dataset() -> dict:
-    datasets = dict()
+    datasets = set()
     for ds in initial_models().keys():
         if initial_models()[ds].get("default", False):
-            datasets[ds] = True
-    return datasets
+            datasets.add(ds)
+    return list(datasets)
 
 
 # ---------------------------------------------
@@ -186,14 +225,14 @@ def all_datasets() -> dict:
 # look for legacy model.ckpt in models directory and offer to
 # normalize its name
 def migrate_models_ckpt():
-    model_path = os.path.join(config.root, Model_dir, Weights_dir)
+    model_path = os.path.join(config.root_dir, Model_dir, Weights_dir)
     if not os.path.exists(os.path.join(model_path, "model.ckpt")):
         return
     new_name = initial_models()["stable-diffusion-1.4"]["file"]
-    print(
+    logger.warning(
         'The Stable Diffusion v4.1 "model.ckpt" is already installed. The name will be changed to {new_name} to avoid confusion.'
     )
-    print(f"model.ckpt => {new_name}")
+    logger.warning(f"model.ckpt => {new_name}")
     os.replace(
         os.path.join(model_path, "model.ckpt"), os.path.join(model_path, new_name)
     )
@@ -206,7 +245,7 @@ def download_weight_datasets(
     migrate_models_ckpt()
     successful = dict()
     for mod in models:
-        print(f"Downloading {mod}:")
+        logger.info(f"Downloading {mod}:")
         successful[mod] = _download_repo_or_file(
             initial_models()[mod], access_token, precision=precision
         )
@@ -227,11 +266,10 @@ def _download_repo_or_file(
             )
     return path
 
-
 def _download_ckpt_weights(mconfig: DictConfig, access_token: str) -> Path:
     repo_id = mconfig["repo_id"]
     filename = mconfig["file"]
-    cache_dir = os.path.join(config.root, Model_dir, Weights_dir)
+    cache_dir = os.path.join(config.root_dir, Model_dir, Weights_dir)
     return hf_download_with_resume(
         repo_id=repo_id,
         model_dir=cache_dir,
@@ -244,6 +282,9 @@ def _download_ckpt_weights(mconfig: DictConfig, access_token: str) -> Path:
 def download_from_hf(
     model_class: object, model_name: str, **kwargs
 ):
+    logger = InvokeAILogger.getLogger('InvokeAI')
+    logger.addFilter(lambda x: 'fp16 is not a valid' not in x.getMessage())
+    
     path = config.cache_dir
     model = model_class.from_pretrained(
         model_name,
@@ -275,10 +316,10 @@ def _download_diffusion_weights(
                 **extra_args,
             )
         except OSError as e:
-            if str(e).startswith("fp16 is not a valid"):
+            if 'Revision Not Found' in str(e):
                 pass
             else:
-                print(f"An unexpected error occurred while downloading the model: {e})")
+                logger.error(str(e))
         if path:
             break
     return path
@@ -286,9 +327,13 @@ def _download_diffusion_weights(
 
 # ---------------------------------------------
 def hf_download_with_resume(
-    repo_id: str, model_dir: str, model_name: str, access_token: str = None
+        repo_id: str,
+        model_dir: str,
+        model_name: str,
+        model_dest: Path = None,
+        access_token: str = None,
 ) -> Path:
-    model_dest = Path(os.path.join(model_dir, model_name))
+    model_dest = model_dest or Path(os.path.join(model_dir, model_name))
     os.makedirs(model_dir, exist_ok=True)
 
     url = hf_hub_url(repo_id, model_name)
@@ -308,20 +353,19 @@ def hf_download_with_resume(
     if (
         resp.status_code == 416
     ):  # "range not satisfiable", which means nothing to return
-        print(f"* {model_name}: complete file found. Skipping.")
+        logger.info(f"{model_name}: complete file found. Skipping.")
         return model_dest
+    elif resp.status_code == 404:
+        logger.warning("File not found")
+        return None
     elif resp.status_code != 200:
-        print(f"** An error occurred during downloading {model_name}: {resp.reason}")
+        logger.warning(f"{model_name}: {resp.reason}")
     elif exist_size > 0:
-        print(f"* {model_name}: partial file found. Resuming...")
+        logger.info(f"{model_name}: partial file found. Resuming...")
     else:
-        print(f"* {model_name}: Downloading...")
+        logger.info(f"{model_name}: Downloading...")
 
     try:
-        if total < 2000:
-            print(f"*** ERROR DOWNLOADING {model_name}: {resp.text}")
-            return None
-
         with open(model_dest, open_mode) as file, tqdm(
             desc=model_name,
             initial=exist_size,
@@ -334,7 +378,7 @@ def hf_download_with_resume(
                 size = file.write(data)
                 bar.update(size)
     except Exception as e:
-        print(f"An error occurred while downloading {model_name}: {str(e)}")
+        logger.error(f"An error occurred while downloading {model_name}: {str(e)}")
         return None
     return model_dest
 
@@ -359,8 +403,8 @@ def update_config_file(successfully_downloaded: dict, config_file: Path):
     try:
         backup = None
         if os.path.exists(config_file):
-            print(
-                f"** {config_file.name} exists. Renaming to {config_file.stem}.yaml.orig"
+            logger.warning(
+                f"{config_file.name} exists. Renaming to {config_file.stem}.yaml.orig"
             )
             backup = config_file.with_suffix(".yaml.orig")
             ## Ugh. Windows is unable to overwrite an existing backup file, raises a WinError 183
@@ -377,16 +421,16 @@ def update_config_file(successfully_downloaded: dict, config_file: Path):
                 new_config.write(tmp.read())
 
     except Exception as e:
-        print(f"**Error creating config file {config_file}: {str(e)} **")
+        logger.error(f"Error creating config file {config_file}: {str(e)}")
         if backup is not None:
-            print("restoring previous config file")
+            logger.info("restoring previous config file")
             ## workaround, for WinError 183, see above
             if sys.platform == "win32" and config_file.is_file():
                 config_file.unlink()
             backup.rename(config_file)
         return
-
-    print(f"Successfully created new configuration file {config_file}")
+    
+    logger.info(f"Successfully created new configuration file {config_file}")
 
 
 # ---------------------------------------------
@@ -420,7 +464,7 @@ def new_config_file_contents(
             stanza["height"] = mod["height"]
         if "file" in mod:
             stanza["weights"] = os.path.relpath(
-                successfully_downloaded[model], start=config.root
+                successfully_downloaded[model], start=config.root_dir
             )
             stanza["config"] = os.path.normpath(
                 os.path.join(sd_configs(), mod["config"])
@@ -453,14 +497,14 @@ def delete_weights(model_name: str, conf_stanza: dict):
     if re.match("/VAE/", conf_stanza.get("config")):
         return
 
-    print(
-        f"\n** The checkpoint version of {model_name} is superseded by the diffusers version. Deleting the original file {weights}?"
+    logger.warning(
+        f"\nThe checkpoint version of {model_name} is superseded by the diffusers version. Deleting the original file {weights}?"
     )
 
     weights = Path(weights)
     if not weights.is_absolute():
-        weights = Path(config.root) / weights
+        weights = config.root_dir / weights
         try:
             weights.unlink()
         except OSError as e:
-            print(str(e))
+            logger.error(str(e))

invokeai/backend/model_management/convert_ckpt_to_diffusers.py

@@ -1093,6 +1093,8 @@ def load_pipeline_from_original_stable_diffusion_ckpt(
     :param vae_path: Path to a checkpoint VAE that will be converted into diffusers and loaded into the pipeline.
     """
     config = InvokeAIAppConfig.get_config()
+    cache_dir = config.cache_dir
+
     with warnings.catch_warnings():
         warnings.simplefilter("ignore")
         verbosity = dlogging.get_verbosity()
@@ -1105,7 +1107,6 @@ def load_pipeline_from_original_stable_diffusion_ckpt(
         else:
             checkpoint = load_file(checkpoint_path)
 
-        cache_dir = config.cache_dir
         pipeline_class = (
             StableDiffusionGeneratorPipeline
             if return_generator_pipeline
@@ -1297,7 +1298,7 @@ def load_pipeline_from_original_stable_diffusion_ckpt(
             )
             safety_checker = StableDiffusionSafetyChecker.from_pretrained(
                 "CompVis/stable-diffusion-safety-checker",
-                cache_dir=config.cache_dir,
+                cache_dir=cache_dir,
             )
             feature_extractor = AutoFeatureExtractor.from_pretrained(
                 "CompVis/stable-diffusion-safety-checker", cache_dir=cache_dir

invokeai/backend/model_management/model_manager.py

@@ -11,14 +11,16 @@ import gc
 import hashlib
 import os
 import re
+import shutil
 import sys
 import textwrap
 import time
+import traceback
 import warnings
 from enum import Enum, auto
 from pathlib import Path
 from shutil import move, rmtree
-from typing import Any, Optional, Union, Callable, types
+from typing import Any, Optional, Union, Callable, Dict, List, types
 
 import safetensors
 import safetensors.torch
@@ -48,6 +50,10 @@ from ..stable_diffusion import (
     StableDiffusionGeneratorPipeline,
 )
 from invokeai.app.services.config import InvokeAIAppConfig
+from ..install.model_install_backend import (
+    Dataset_path,
+    hf_download_with_resume,
+)
 from ..util import CUDA_DEVICE, ask_user, download_with_resume
 
 class SDLegacyType(Enum):
@@ -314,7 +320,7 @@ class ModelManager(object):
         models = {}
         for name in sorted(self.config, key=str.casefold):
             stanza = self.config[name]
-
+            
             # don't include VAEs in listing (legacy style)
             if "config" in stanza and "/VAE/" in stanza["config"]:
                 continue
@@ -526,7 +532,8 @@ class ModelManager(object):
                     **fp_args,
                 )
             except OSError as e:
-                if str(e).startswith("fp16 is not a valid"):
+                if str(e).startswith("fp16 is not a valid") or \
+                   'Invalid rev id: fp16' in str(e):
                     pass
                 else:
                     self.logger.error(
@@ -773,7 +780,7 @@ class ModelManager(object):
 
         """
         model_path: Path = None
-        thing = path_url_or_repo  # to save typing
+        thing = str(path_url_or_repo)  # to save typing
 
         self.logger.info(f"Probing {thing} for import")
 
@@ -813,7 +820,9 @@ class ModelManager(object):
                     Path(thing).rglob("*.safetensors")
                 ):
                     if model_name := self.heuristic_import(
-                        str(m), commit_to_conf=commit_to_conf
+                            str(m),
+                            commit_to_conf=commit_to_conf,
+                            config_file_callback=config_file_callback,
                     ):
                         self.logger.info(f"{model_name} successfully imported")
                 return model_name
@@ -859,7 +868,7 @@ class ModelManager(object):
                     model_config_file = self.globals.legacy_conf_path / "v1-inference.yaml"
                 elif model_type == SDLegacyType.V1_INPAINT:
                     self.logger.debug("SD-v1 inpainting model detected")
-                    model_config_file = self.globals.legacy_conf_path / "v1-inpainting-inference.yaml",
+                    model_config_file = self.globals.legacy_conf_path / "v1-inpainting-inference.yaml"
                 elif model_type == SDLegacyType.V2_v:
                     self.logger.debug("SD-v2-v model detected")
                     model_config_file = self.globals.legacy_conf_path / "v2-inference-v.yaml"
@@ -868,14 +877,12 @@ class ModelManager(object):
                     model_config_file = self.globals.legacy_conf_path / "v2-inference.yaml"
                 elif model_type == SDLegacyType.V2:
                     self.logger.warning(
-                        f"{thing} is a V2 checkpoint file, but its parameterization cannot be determined. Please provide configuration file path."
+                        f"{thing} is a V2 checkpoint file, but its parameterization cannot be determined."
                     )
-                    return
                 else:
                     self.logger.warning(
-                        f"{thing} is a legacy checkpoint file but not a known Stable Diffusion model. Please provide configuration file path."
+                        f"{thing} is a legacy checkpoint file but not a known Stable Diffusion model."
                     )
-                    return
 
         if not model_config_file and config_file_callback:
             model_config_file = config_file_callback(model_path)
@@ -932,34 +939,35 @@ class ModelManager(object):
 
         from . import convert_ckpt_to_diffusers
 
-        if diffusers_path.exists():
-            self.logger.error(
-                f"The path {str(diffusers_path)} already exists. Please move or remove it and try again."
-            )
-            return
-
         model_name = model_name or diffusers_path.name
         model_description = model_description or f"Converted version of {model_name}"
-        self.logger.debug(f"Converting {model_name} to diffusers (30-60s)")
+        
         try:
-            # By passing the specified VAE to the conversion function, the autoencoder
-            # will be built into the model rather than tacked on afterward via the config file
-            vae_model = None
-            if vae:
-                vae_model = self._load_vae(vae)
-                vae_path = None
-            convert_ckpt_to_diffusers(
-                ckpt_path,
-                diffusers_path,
-                extract_ema=True,
-                original_config_file=original_config_file,
-                vae=vae_model,
-                vae_path=vae_path,
-                scan_needed=scan_needed,
-            )
-            self.logger.debug(
-                f"Success. Converted model is now located at {str(diffusers_path)}"
-            )
+            if diffusers_path.exists():
+                self.logger.error(
+                    f"The path {str(diffusers_path)} already exists. Installing previously-converted path."
+                )
+            else:
+                self.logger.debug(f"Converting {model_name} to diffusers (30-60s)")
+
+                # By passing the specified VAE to the conversion function, the autoencoder
+                # will be built into the model rather than tacked on afterward via the config file
+                vae_model = None
+                if vae:
+                    vae_model = self._load_vae(vae)
+                    vae_path = None
+                convert_ckpt_to_diffusers(
+                    ckpt_path,
+                    diffusers_path,
+                    extract_ema=True,
+                    original_config_file=original_config_file,
+                    vae=vae_model,
+                    vae_path=vae_path,
+                    scan_needed=scan_needed,
+                )
+                self.logger.debug(
+                    f"Success. Converted model is now located at {str(diffusers_path)}"
+                )
             self.logger.debug(f"Writing new config file entry for {model_name}")
             new_config = dict(
                 path=str(diffusers_path),
@@ -971,9 +979,10 @@ class ModelManager(object):
             self.add_model(model_name, new_config, True)
             if commit_to_conf:
                 self.commit(commit_to_conf)
-            self.logger.debug("Conversion succeeded")
+            self.logger.debug(f"Model {model_name} installed")
         except Exception as e:
             self.logger.warning(f"Conversion failed: {str(e)}")
+            self.logger.warning(traceback.format_exc())
             self.logger.warning(
                 "If you are trying to convert an inpainting or 2.X model, please indicate the correct config file (e.g. v1-inpainting-inference.yaml)"
             )
@@ -1314,3 +1323,185 @@ class ModelManager(object):
         return (
             os.getenv("HF_HOME") is not None or os.getenv("XDG_CACHE_HOME") is not None
         )
+
+    def list_lora_models(self)->Dict[str,bool]:
+        '''Return a dict of installed lora models; key is either the shortname
+        defined in INITIAL_MODELS, or the basename of the file in the LoRA
+        directory. Value is True if installed'''
+
+        models = OmegaConf.load(Dataset_path).get('lora') or {}
+        installed_models = {x: False for x in models.keys()}
+        
+        dir = self.globals.lora_path
+        installed_models = dict()
+        for root, dirs, files in os.walk(dir):
+            for name in files:
+                if Path(name).suffix not in ['.safetensors','.ckpt','.pt','.bin']:
+                    continue
+                if name == 'pytorch_lora_weights.bin':
+                    name = Path(root,name).parent.stem #Path(root,name).stem
+                else:
+                    name = Path(name).stem
+                installed_models.update({name: True})
+
+        return installed_models
+    
+    def install_lora_models(self, model_names: list[str], access_token:str=None):
+        '''Download list of LoRA/LyCORIS models'''
+        
+        short_names = OmegaConf.load(Dataset_path).get('lora') or {}
+        for name in model_names:
+            name = short_names.get(name) or name
+
+            # HuggingFace style LoRA
+            if re.match(r"^[\w.+-]+/([\w.+-]+)$", name):
+                self.logger.info(f'Downloading LoRA/LyCORIS model {name}')
+                _,dest_dir = name.split("/")
+                
+                hf_download_with_resume(
+                    repo_id = name,
+                    model_dir = self.globals.lora_path / dest_dir,
+                    model_name = 'pytorch_lora_weights.bin',
+                    access_token = access_token,
+                )
+
+            elif name.startswith(("http:", "https:", "ftp:")):
+                download_with_resume(name, self.globals.lora_path)
+
+            else:
+                self.logger.error(f"Unknown repo_id or URL: {name}")
+    
+    def delete_lora_models(self, model_names: List[str]):
+        '''Remove the list of lora models'''
+        for name in model_names:
+            file_or_directory = self.globals.lora_path / name
+            if file_or_directory.is_dir():
+                self.logger.info(f'Purging LoRA/LyCORIS {name}')
+                shutil.rmtree(str(file_or_directory))
+            else:
+                for path in self.globals.lora_path.glob(f'{name}.*'):
+                    self.logger.info(f'Purging LoRA/LyCORIS {name}')
+                    path.unlink()
+                
+    def list_ti_models(self)->Dict[str,bool]:
+        '''Return a dict of installed textual models; key is either the shortname
+        defined in INITIAL_MODELS, or the basename of the file in the LoRA
+        directory. Value is True if installed'''
+
+        models = OmegaConf.load(Dataset_path).get('textual_inversion') or {}
+        installed_models = {x: False for x in models.keys()}
+        
+        dir = self.globals.embedding_path
+        for root, dirs, files in os.walk(dir):
+            for name in files:
+                if not Path(name).suffix in ['.bin','.pt','.ckpt','.safetensors']:
+                    continue
+                if name == 'learned_embeds.bin':
+                    name = Path(root,name).parent.stem #Path(root,name).stem
+                else:
+                    name = Path(name).stem
+                installed_models.update({name: True})
+        return installed_models
+
+    def install_ti_models(self, model_names: list[str], access_token: str=None):
+        '''Download list of textual inversion embeddings'''
+
+        short_names = OmegaConf.load(Dataset_path).get('textual_inversion') or {}
+        for name in model_names:
+            name = short_names.get(name) or name
+            
+            if re.match(r"^[\w.+-]+/([\w.+-]+)$", name):
+                self.logger.info(f'Downloading Textual Inversion embedding {name}')
+                _,dest_dir = name.split("/")
+                hf_download_with_resume(
+                    repo_id = name,
+                    model_dir = self.globals.embedding_path / dest_dir,
+                    model_name = 'learned_embeds.bin',
+                    access_token = access_token
+                )
+            elif name.startswith(('http:','https:','ftp:')):
+                download_with_resume(name, self.globals.embedding_path)
+            else:
+                self.logger.error(f'{name} does not look like either a HuggingFace repo_id or a downloadable URL')
+    
+    def delete_ti_models(self, model_names: list[str]):
+        '''Remove TI embeddings from disk'''
+        for name in model_names:
+            file_or_directory = self.globals.embedding_path / name
+            if file_or_directory.is_dir():
+                self.logger.info(f'Purging textual inversion embedding {name}')
+                shutil.rmtree(str(file_or_directory))
+            else:
+                for path in self.globals.embedding_path.glob(f'{name}.*'):
+                    self.logger.info(f'Purging textual inversion embedding {name}')
+                    path.unlink()
+
+    def list_controlnet_models(self)->Dict[str,bool]:
+        '''Return a dict of installed controlnet models; key is repo_id or short name
+        of model (defined in INITIAL_MODELS), and value is True if installed'''
+
+        cn_models = OmegaConf.load(Dataset_path).get('controlnet') or {}
+        installed_models = {x: False for x in cn_models.keys()}
+        
+        cn_dir = self.globals.controlnet_path
+        for root, dirs, files in os.walk(cn_dir):
+            for name in dirs:
+                if Path(root, name, '.download_complete').exists():
+                    installed_models.update({name.replace('--','/'): True})
+        return installed_models
+
+    def install_controlnet_models(self, model_names: list[str], access_token: str=None):
+        '''Download list of controlnet models; provide either repo_id or short name listed in INITIAL_MODELS.yaml'''
+        short_names = OmegaConf.load(Dataset_path).get('controlnet') or {}
+        dest_dir = self.globals.controlnet_path
+        dest_dir.mkdir(parents=True,exist_ok=True)
+        
+        # The model file may be fp32 or fp16, and may be either a
+        # .bin file or a .safetensors. We try each until we get one,
+        # preferring 'fp16' if using half precision, and preferring
+        # safetensors over over bin.
+        precisions = ['.fp16',''] if self.precision=='float16' else ['']
+        formats = ['.safetensors','.bin']
+        possible_filenames = list()
+        for p in precisions:
+            for f in formats:
+                possible_filenames.append(Path(f'diffusion_pytorch_model{p}{f}'))
+        
+        for directory_name in model_names:
+            repo_id = short_names.get(directory_name) or directory_name
+            safe_name = directory_name.replace('/','--')
+            self.logger.info(f'Downloading ControlNet model {directory_name} ({repo_id})')
+            hf_download_with_resume(
+                repo_id = repo_id,
+                model_dir = dest_dir / safe_name,
+                model_name = 'config.json',
+                access_token = access_token
+            )
+
+            path = None
+            for filename in possible_filenames:
+                suffix = filename.suffix
+                dest_filename = Path(f'diffusion_pytorch_model{suffix}')
+                self.logger.info(f'Checking availability of {directory_name}/{filename}...')
+                path = hf_download_with_resume(
+                    repo_id = repo_id,
+                    model_dir = dest_dir / safe_name,
+                    model_name = str(filename),
+                    access_token = access_token,
+                    model_dest = Path(dest_dir, safe_name, dest_filename),
+                )
+                if path:
+                    (path.parent / '.download_complete').touch()
+                    break
+
+    def delete_controlnet_models(self, model_names: List[str]):
+        '''Remove the list of controlnet models'''
+        for name in model_names:
+            safe_name = name.replace('/','--')
+            directory = self.globals.controlnet_path / safe_name
+            if directory.exists():
+                self.logger.info(f'Purging controlnet model {name}')
+                shutil.rmtree(str(directory))
+
+
+        

invokeai/backend/util/logging.py

@@ -30,7 +30,6 @@ import invokeai.backend.util.logging as IAILogger
 IAILogger.debug('this is a debugging message')
 """
 
-import logging
 import logging.handlers
 import socket
 import urllib.parse
@@ -195,7 +194,6 @@ class InvokeAILogger(object):
     @classmethod
     def getLoggers(cls, config: InvokeAIAppConfig) -> list[logging.Handler]:
         handler_strs = config.log_handlers
-        print(f'handler_strs={handler_strs}')
         handlers = list()
         for handler in handler_strs:
             handler_name,*args = handler.split('=',2)

invokeai/backend/util/util.py

@@ -322,8 +322,8 @@ def download_with_resume(url: str, dest: Path, access_token: str = None) -> Path
         logger.warning("corrupt existing file found. re-downloading")
         os.remove(dest)
         exist_size = 0
-
-    if resp.status_code == 416 or exist_size == content_length:
+        
+    if resp.status_code == 416 or (content_length > 0 and exist_size == content_length):
         logger.warning(f"{dest}: complete file found. Skipping.")
         return dest
     elif resp.status_code == 206 or exist_size > 0:
@@ -331,7 +331,7 @@ def download_with_resume(url: str, dest: Path, access_token: str = None) -> Path
     elif resp.status_code != 200:
         logger.error(f"An error occurred during downloading {dest}: {resp.reason}")
     else:
-        logger.error(f"{dest}: Downloading...")
+        logger.info(f"{dest}: Downloading...")
 
     try:
         if content_length < 2000: