From 1632ac6b9f30d8efdb2fa09ebbd41fd0f58fae3b Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Tue, 30 May 2023 13:49:43 -0400 Subject: [PATCH] add controlnet model downloading --- invokeai/app/services/config.py | 9 +- .../backend/config/model_install_backend.py | 140 ++++++++++--- invokeai/frontend/install/model_install.py | 184 +++++++++--------- 3 files changed, 206 insertions(+), 127 deletions(-) diff --git a/invokeai/app/services/config.py b/invokeai/app/services/config.py index 75ae034235..c4903b918a 100644 --- a/invokeai/app/services/config.py +++ b/invokeai/app/services/config.py @@ -359,7 +359,7 @@ setting environment variables INVOKEAI_. conf_path : Path = Field(default='configs/models.yaml', description='Path to models definition file', category='Paths') embedding_dir : Path = Field(default='embeddings', description='Path to InvokeAI textual inversion aembeddings directory', category='Paths') gfpgan_model_dir : Path = Field(default="./models/gfpgan/GFPGANv1.4.pth", description='Path to GFPGAN models directory.', category='Paths') - controlnet_dir : Path = Field(default="controlnet", description='Path to directory of ControlNet models.', category='Paths') + controlnet_dir : Path = Field(default="controlnets", description='Path to directory of ControlNet models.', category='Paths') legacy_conf_dir : Path = Field(default='configs/stable-diffusion', description='Path to directory of legacy checkpoint config files', category='Paths') lora_dir : Path = Field(default='loras', description='Path to InvokeAI LoRA model directory', category='Paths') outdir : Path = Field(default='outputs', description='Default folder for output images', category='Paths') @@ -417,6 +417,13 @@ setting environment variables INVOKEAI_. def _resolve(self,partial_path:Path)->Path: return (self.root_path / partial_path).resolve() + @property + def init_file_path(self)->Path: + ''' + Path to invokeai.yaml + ''' + return self._resolve(INIT_FILE) + @property def output_path(self)->Path: ''' diff --git a/invokeai/backend/config/model_install_backend.py b/invokeai/backend/config/model_install_backend.py index 644ecea0fc..28bc42cb4b 100644 --- a/invokeai/backend/config/model_install_backend.py +++ b/invokeai/backend/config/model_install_backend.py @@ -8,11 +8,11 @@ import sys import warnings from pathlib import Path from tempfile import TemporaryFile -from typing import List +from typing import List, Dict 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 @@ -49,7 +49,6 @@ Config_preamble = """ def default_config_file(): - print(config.root_dir) return config.model_conf_path def sd_configs(): @@ -62,23 +61,35 @@ def initial_models(): 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, + install_initial_models: List[str] = None, + remove_models: List[str] = None, + install_cn_models: List[str] = None, + remove_cn_models: List[str] = None, + cn_model_map: Dict[str,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, ): """ 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") - model_manager = ModelManager(OmegaConf.load(config_file_path)['diffusers'], precision=precision) + install_controlnet_models( + install_cn_models, + short_name_map = cn_model_map, + precision=precision, + access_token=access_token, + ) + delete_controlnet_models(remove_cn_models) + + model_manager = ModelManager(OmegaConf.load(config_file_path), precision=precision) if remove_models and len(remove_models) > 0: print("== DELETING UNCHECKED STARTER MODELS ==") @@ -120,18 +131,20 @@ 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_dir, 'invokeai.init') - replacement = Path(config.root_dir, 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) + +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) + 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) # ------------------------------------- @@ -227,6 +240,68 @@ def _download_ckpt_weights(mconfig: DictConfig, access_token: str) -> Path: ) +# --------------------------------------------- +def install_controlnet_models( + short_names: List[str], + short_name_map: Dict[str,str], + precision: str='float16', + access_token: str = None, +): + ''' + Download list of controlnet models, using their HuggingFace + repo_ids. + ''' + dest_dir = config.controlnet_path + if not dest_dir.exists(): + dest_dir.mkdir(parents=True,exist_ok=False) + + # 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 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 short_names: + repo_id = short_name_map[directory_name] + safe_name = directory_name.replace('/','--') + print(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}') + print(f'Probing {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(short_names: List[str]): + for name in short_names: + safe_name = name.replace('/','--') + directory = config.controlnet_path / safe_name + if directory.exists(): + print(f'Purging controlnet model {name}') + shutil.rmtree(str(directory)) + # --------------------------------------------- def download_from_hf( model_class: object, model_name: str, **kwargs @@ -273,9 +348,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) @@ -297,18 +376,17 @@ def hf_download_with_resume( ): # "range not satisfiable", which means nothing to return print(f"* {model_name}: complete file found. Skipping.") return model_dest + elif resp.status_code == 404: + print("** File not found") + return None elif resp.status_code != 200: - print(f"** An error occurred during downloading {model_name}: {resp.reason}") + print(f"** Warning: {model_name}: {resp.reason}") elif exist_size > 0: print(f"* {model_name}: partial file found. Resuming...") else: print(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, diff --git a/invokeai/frontend/install/model_install.py b/invokeai/frontend/install/model_install.py index 59b2075b28..1a9f8117b4 100644 --- a/invokeai/frontend/install/model_install.py +++ b/invokeai/frontend/install/model_install.py @@ -43,7 +43,7 @@ from invokeai.app.services.config import get_invokeai_config # minimum size for the UI MIN_COLS = 120 -MIN_LINES = 45 +MIN_LINES = 50 config = get_invokeai_config() @@ -53,16 +53,16 @@ class addModelsForm(npyscreen.FormMultiPage): def __init__(self, parentApp, name, multipage=False, *args, **keywords): self.multipage = multipage + self.initial_models = OmegaConf.load(Dataset_path)['diffusers'] self.control_net_models = OmegaConf.load(Dataset_path)['controlnet'] self.installed_cn_models = self._get_installed_cn_models() + self._add_additional_cn_models(self.control_net_models,self.installed_cn_models) + try: self.existing_models = OmegaConf.load(default_config_file()) except: self.existing_models = dict() - # self.starter_model_list = [ - # x for x in list(self.initial_models.keys()) if x not in self.existing_models - # ] self.starter_model_list = list(self.initial_models.keys()) self.installed_models = dict() super().__init__(parentApp=parentApp, name=name, *args, **keywords) @@ -95,40 +95,6 @@ class addModelsForm(npyscreen.FormMultiPage): color="CAUTION", ) self.nextrely += 1 - # if len(self.installed_models) > 0: - # self.add_widget_intelligent( - # CenteredTitleText, - # name="== INSTALLED STARTER MODELS ==", - # editable=False, - # color="CONTROL", - # ) - # self.nextrely -= 1 - # self.add_widget_intelligent( - # CenteredTitleText, - # name="Currently installed starter models. Uncheck to delete:", - # editable=False, - # labelColor="CAUTION", - # ) - # self.nextrely -= 1 - # columns = self._get_columns() - # self.previously_installed_models = self.add_widget_intelligent( - # MultiSelectColumns, - # columns=columns, - # values=self.installed_models, - # value=[x for x in range(0, len(self.installed_models))], - # max_height=1 + len(self.installed_models) // columns, - # relx=4, - # slow_scroll=True, - # scroll_exit=True, - # ) - # self.purge_deleted = self.add_widget_intelligent( - # npyscreen.Checkbox, - # name="Purge deleted models from disk", - # value=False, - # scroll_exit=True, - # relx=4, - # ) - # self.nextrely += 1 if len(self.starter_model_list) > 0: self.add_widget_intelligent( CenteredTitleText, @@ -161,35 +127,14 @@ class addModelsForm(npyscreen.FormMultiPage): relx=4, scroll_exit=True, ) - self.add_widget_intelligent( - CenteredTitleText, - name="== CONTROLNET MODELS ==", - editable=False, - color="CONTROL", - ) - columns=6 - self.cn_models_selected = self.add_widget_intelligent( - MultiSelectColumns, - columns=columns, - name="Install ControlNet Models", - values=cn_model_list, - value=[ - cn_model_list.index(x) - for x in cn_model_list - if x in self.installed_cn_models - ], - max_height=len(cn_model_list)//columns + 1, - relx=4, - scroll_exit=True, - ) - self.nextrely += 1 self.purge_deleted = self.add_widget_intelligent( npyscreen.Checkbox, - name="Purge unchecked models from disk", + name="Purge unchecked diffusers models from disk", value=False, scroll_exit=True, relx=4, ) + self.nextrely += 1 self.add_widget_intelligent( CenteredTitleText, name="== IMPORT LOCAL AND REMOTE MODELS ==", @@ -211,7 +156,7 @@ class addModelsForm(npyscreen.FormMultiPage): ) self.nextrely -= 1 self.import_model_paths = self.add_widget_intelligent( - TextBox, max_height=7, scroll_exit=True, editable=True, relx=4 + TextBox, max_height=4, scroll_exit=True, editable=True, relx=4 ) self.nextrely += 1 self.show_directory_fields = self.add_widget_intelligent( @@ -236,6 +181,47 @@ class addModelsForm(npyscreen.FormMultiPage): relx=4, scroll_exit=True, ) + self.add_widget_intelligent( + CenteredTitleText, + name="== CONTROLNET MODELS ==", + editable=False, + color="CONTROL", + ) + self.nextrely -= 1 + self.add_widget_intelligent( + CenteredTitleText, + name="Select the desired ControlNet models. Unchecked models will be purged from disk.", + editable=False, + labelColor="CAUTION", + ) + columns=6 + self.cn_models_selected = self.add_widget_intelligent( + MultiSelectColumns, + columns=columns, + name="Install ControlNet Models", + values=cn_model_list, + value=[ + cn_model_list.index(x) + for x in cn_model_list + if x in self.installed_cn_models + ], + max_height=len(cn_model_list)//columns + 1, + relx=4, + scroll_exit=True, + ) + self.nextrely += 1 + self.add_widget_intelligent( + npyscreen.TitleFixedText, + name='Additional ControlNet HuggingFace repo_ids to install (space separated):', + relx=4, + color='CONTROL', + editable=False, + scroll_exit=True + ) + self.nextrely -= 1 + self.additional_controlnet_ids = self.add_widget_intelligent( + TextBox, max_height=2, scroll_exit=True, editable=True, relx=4 + ) self.cancel = self.add_widget_intelligent( npyscreen.ButtonPress, name="CANCEL", @@ -300,20 +286,22 @@ class addModelsForm(npyscreen.FormMultiPage): ] def _get_installed_cn_models(self)->list[str]: - with open('log.txt','w') as file: - cn_dir = config.controlnet_path - file.write(f'cn_dir={cn_dir}\n') - installed_cn_models = set() - for root, dirs, files in os.walk(cn_dir): - for name in dirs: - file.write(f'{root}/{name}/config.json\n') - if Path(root, name, 'config.json').exists(): - installed_cn_models.add(name) - inverse_dict = {name.split('/')[1]: key for key, name in self.control_net_models.items()} - file.write(f'inverse={inverse_dict}') - return [inverse_dict[x] for x in installed_cn_models] - + cn_dir = config.controlnet_path + installed_cn_models = set() + for root, dirs, files in os.walk(cn_dir): + for name in dirs: + if Path(root, name, '.download_complete').exists(): + installed_cn_models.add(name.replace('--','/')) + return installed_cn_models + def _add_additional_cn_models(self, known_models: dict, installed_models: set): + for i in installed_models: + if i in known_models: + continue + # translate from name to repo_id + repo_id = i.replace('--','/') + known_models.update({i: repo_id}) + def _get_columns(self) -> int: window_width, window_height = get_terminal_size() cols = ( @@ -374,15 +362,20 @@ class addModelsForm(npyscreen.FormMultiPage): selections.install_models = [x for x in starter_models if x not in self.existing_models] selections.remove_models = [x for x in self.starter_model_list if x in self.existing_models and x not in starter_models] - selections.install_cn_models = [self.control_net_models[self.cn_models_selected.values[x]] + selections.control_net_map = self.control_net_models + selections.install_cn_models = [self.cn_models_selected.values[x] for x in self.cn_models_selected.value if self.cn_models_selected.values[x] not in self.installed_cn_models ] - selections.remove_cn_models = [self.control_net_models[x] + selections.remove_cn_models = [x for x in self.cn_models_selected.values if x in self.installed_cn_models and self.cn_models_selected.values.index(x) not in self.cn_models_selected.value ] + if (additional_cns := self.additional_controlnet_ids.value.split()): + valid_cns = [x for x in additional_cns if '/' in x] + selections.install_cn_models.extend(valid_cns) + selections.control_net_map.update({x: x for x in valid_cns}) # load directory and whether to scan on startup if self.show_directory_fields.value: @@ -406,6 +399,7 @@ class AddModelApplication(npyscreen.NPSAppManaged): purge_deleted_models=False, install_cn_models = None, remove_cn_models = None, + control_net_map = None, scan_directory=None, autoscan_on_startup=None, import_model_paths=None, @@ -425,24 +419,24 @@ def process_and_execute(opt: Namespace, selections: Namespace): directory_to_scan = selections.scan_directory scan_at_startup = selections.autoscan_on_startup potential_models_to_install = selections.import_model_paths - - print('NOT INSTALLING MODELS DURING DEBUGGING') - print('models to install:',models_to_install) - print('models to remove:',models_to_remove) - print('CN models to install:',selections.install_cn_models) - print('CN models to remove:',selections.remove_cn_models) - # install_requested_models( - # install_initial_models=models_to_install, - # remove_models=models_to_remove, - # scan_directory=Path(directory_to_scan) if directory_to_scan else None, - # external_models=potential_models_to_install, - # scan_at_startup=scan_at_startup, - # precision="float32" - # if opt.full_precision - # else choose_precision(torch.device(choose_torch_device())), - # purge_deleted=selections.purge_deleted_models, - # config_file_path=Path(opt.config_file) if opt.config_file else None, - # ) + print(f'selections.install_cn_models={selections.install_cn_models}') + print(f'selections.remove_cn_models={selections.remove_cn_models}') + print(f'selections.cn_model_map={selections.control_net_map}') + install_requested_models( + install_initial_models=models_to_install, + remove_models=models_to_remove, + install_cn_models=selections.install_cn_models, + remove_cn_models=selections.remove_cn_models, + cn_model_map=selections.control_net_map, + scan_directory=Path(directory_to_scan) if directory_to_scan else None, + external_models=potential_models_to_install, + scan_at_startup=scan_at_startup, + precision="float32" + if opt.full_precision + else choose_precision(torch.device(choose_torch_device())), + purge_deleted=selections.purge_deleted_models, + config_file_path=Path(opt.config_file) if opt.config_file else None, + ) # --------------------------------------------------------