From e87a2fe14b0d4670b7c063f3689e139787ef3d72 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Wed, 15 Feb 2023 01:07:39 -0500 Subject: [PATCH] model installer frontend done - needs to be hooked to backend --- ldm/invoke/config/model_install.py | 585 +++++++----------------- ldm/invoke/config/model_install_util.py | 378 +++++++++++++++ ldm/invoke/config/widgets.py | 27 ++ 3 files changed, 570 insertions(+), 420 deletions(-) create mode 100644 ldm/invoke/config/model_install_util.py diff --git a/ldm/invoke/config/model_install.py b/ldm/invoke/config/model_install.py index 84ba124658..415c99f8f8 100644 --- a/ldm/invoke/config/model_install.py +++ b/ldm/invoke/config/model_install.py @@ -3,79 +3,27 @@ # Before running stable-diffusion on an internet-isolated machine, # run this script from one with internet connectivity. The # two machines must share a common .cache directory. -# -# Coauthor: Kevin Turner http://github.com/keturn -# + import argparse import curses import os -import re -import shutil import sys import traceback -import warnings from argparse import Namespace -from math import ceil -from pathlib import Path -from tempfile import TemporaryFile +from typing import List import npyscreen -import requests -from diffusers import AutoencoderKL -from huggingface_hub import hf_hub_url +import torch from npyscreen import widget from omegaconf import OmegaConf -from omegaconf.dictconfig import DictConfig -from tqdm import tqdm - -import invokeai.configs as configs -from ldm.invoke.devices import choose_precision, choose_torch_device -from ldm.invoke.generator.diffusers_pipeline import StableDiffusionGeneratorPipeline -from ldm.invoke.globals import Globals, global_cache_dir, global_config_dir -from ldm.invoke.config.widgets import MultiSelectColumns - -warnings.filterwarnings("ignore") -import torch - -# --------------------------globals----------------------- -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 = Path(global_config_dir()) / "models.yaml" -SD_Configs = Path(global_config_dir()) / "stable-diffusion" - -Datasets = OmegaConf.load(Dataset_path) - -Config_preamble = """# This file describes the alternative machine learning models -# available to InvokeAI script. -# -# To add a new model, follow the examples below. Each -# model requires a model config file, a weights file, -# and the width and height of the images it -# was trained on. -""" - -# ------------------------------------- -def yes_or_no(prompt: str, default_yes=True): - default = "y" if default_yes else "n" - response = input(f"{prompt} [{default}] ") or default - if default_yes: - return response[0] not in ("n", "N") - 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 Globals.root +from ..devices import choose_precision, choose_torch_device +from ..globals import Globals +from .widgets import MultiSelectColumns, TextBox +from .model_install_util import (Dataset_path, Default_config_file, + default_dataset, download_weight_datasets, + update_config_file, get_root + ) class addModelsForm(npyscreen.FormMultiPageAction): def __init__(self, parentApp, name): @@ -98,38 +46,53 @@ class addModelsForm(npyscreen.FormMultiPageAction): for x in self.starter_model_list if self.initial_models[x].get("recommended", False) ] - previously_installed_models = sorted( + self.installed_models = sorted( [ x for x in list(self.initial_models.keys()) if x in self.existing_models ] ) - if len(previously_installed_models) > 0: - title = self.add_widget_intelligent( - npyscreen.TitleText, - name="Currently installed starter models. Uncheck to delete:", + self.add_widget_intelligent( + npyscreen.FixedText, + value='Use ctrl-N and ctrl-P to move to the ext and

revious fields,', + editable=False, + ) + self.add_widget_intelligent( + npyscreen.FixedText, + value='cursor arrows to make a selection, and space to toggle checkboxes.', + editable=False, + ) + + if len(self.installed_models) > 0: + self.add_widget_intelligent( + npyscreen.TitleFixedText, + name="== INSTALLED STARTER MODELS ==", + value="Currently installed starter models. Uncheck to delete:", + begin_entry_at=2, editable=False, color="CONTROL", ) - self.nextrely -= 1 - columns = 3 + columns = self._get_columns() self.previously_installed_models = self.add_widget_intelligent( MultiSelectColumns, columns=columns, - values=previously_installed_models, - value=[x for x in range(0,len(previously_installed_models))], - max_height=len(previously_installed_models)+1 // columns, + values=self.installed_models, + value=[x for x in range(0,len(self.installed_models))], + max_height=2+len(self.installed_models) // columns, + relx = 4, slow_scroll=True, scroll_exit = True, ) + self.add_widget_intelligent( - npyscreen.TitleText, - name="Select from a starter set of Stable Diffusion models from HuggingFace:", + npyscreen.TitleFixedText, + name="== UNINSTALLED STARTER MODELS ==", + value="Select from a starter set of Stable Diffusion models from HuggingFace:", + begin_entry_at=2, editable=False, color="CONTROL", ) - self.nextrely -= 2 - self.add_widget_intelligent(npyscreen.FixedText, value="", editable=False), + self.nextrely -= 1 self.models_selected = self.add_widget_intelligent( npyscreen.MultiSelect, name="Install Starter Models", @@ -140,39 +103,39 @@ class addModelsForm(npyscreen.FormMultiPageAction): if x in recommended_models ], max_height=len(starter_model_labels) + 1, + relx = 4, scroll_exit=True, ) - for line in [ - 'Import checkpoint/safetensor models from the directory below.', - '(Use to autocomplete)' - ]: - self.add_widget_intelligent( - npyscreen.TitleText, - name=line, - editable=False, - color="CONTROL", - ) - self.nextrely -= 1 + self.add_widget_intelligent( + npyscreen.TitleFixedText, + name='== MODEL IMPORT DIRECTORY ==', + value='Import all models found in this directory ( autocompletes):', + begin_entry_at=2, + editable=False, + color="CONTROL", + ) self.autoload_directory = self.add_widget_intelligent( npyscreen.TitleFilename, name='Directory:', select_dir=True, must_exist=True, use_two_lines=False, - value=os.path.expanduser('~'+'/'), + relx = 4, + labelColor='DANGER', scroll_exit=True, ) - self.autoload_onstartup = self.add_widget_intelligent( + self.autoscan_on_startup = self.add_widget_intelligent( npyscreen.Checkbox, name='Scan this directory each time InvokeAI starts for new models to import.', value=False, + relx = 4, scroll_exit=True, ) self.nextrely += 1 for line in [ - 'In the space below, you may cut and paste URLs, paths to .ckpt/.safetensor files', - 'or HuggingFace diffusers repository names to import.', - '(Use control-V or shift-control-V to paste):' + '== INDIVIDUAL MODELS TO IMPORT ==', + 'Enter list of URLs, paths models or HuggingFace diffusers repository IDs.', + 'Use control-V or shift-control-V to paste:' ]: self.add_widget_intelligent( npyscreen.TitleText, @@ -181,27 +144,33 @@ class addModelsForm(npyscreen.FormMultiPageAction): color="CONTROL", ) self.nextrely -= 1 - self.model_names = self.add_widget_intelligent( - npyscreen.MultiLineEdit, - max_width=75, + self.import_model_paths = self.add_widget_intelligent( + TextBox, max_height=8, scroll_exit=True, - relx=3 + editable=True, + relx=4 ) - self.autoload_onstartup = self.add_widget_intelligent( + self.nextrely += 2 + self.convert_models = self.add_widget_intelligent( npyscreen.TitleSelectOne, - name='Keep files in original format, or convert .ckpt/.safetensors into fast-loading diffusers models:', - values=['Original format','Convert to diffusers format'], + name='== CONVERT IMPORTED MODELS INTO DIFFUSERS==', + values=['Keep original format','Convert to diffusers'], value=0, + begin_entry_at=4, scroll_exit=True, ) - self.find_next_editable() -# self.set_editing(self.models_selected) -# self.display() -# self.models_selected.editing=True -# self.models_selected.edit() + + def resize(self): + super().resize() + self.models_selected.values = self._get_starter_model_labels() + # thought this would dynamically resize the widget, but no luck + # self.previously_installed_models.columns = self._get_columns() + # self.previously_installed_models.max_height = 2+len(self.installed_models) // self._get_columns() + # self.previously_installed_models.make_contained_widgets() + # self.previously_installed_models.display() - def _get_starter_model_labels(self): + def _get_starter_model_labels(self)->List[str]: window_height, window_width = curses.initscr().getmaxyx() label_width = 25 checkbox_width = 4 @@ -217,17 +186,89 @@ class addModelsForm(npyscreen.FormMultiPageAction): f"%-{label_width}s %s" % (names[x], descriptions[x]) for x in range(0,len(im)) ] + def _get_columns(self)->int: + window_height, window_width = curses.initscr().getmaxyx() + return 4 if window_width > 240 else 3 if window_width>160 else 2 if window_width>80 else 1 + def on_ok(self): - self.parentApp.setNextForm(None) + self.parentApp.setNextForm('MONITOR_OUTPUT') self.editing = False - self.parentApp.selected_models = [ - self.starter_model_list[x] for x in self.models_selected.value - ] - npyscreen.notify(f"Installing selected {self.parentApp.selected_models}") + self.parentApp.user_cancelled = False + self.marshall_arguments() def on_cancel(self): self.parentApp.setNextForm(None) - self.parentApp.selected_models = None + self.ParentApp.user_cancelled = True + self.editing = False + + def marshall_arguments(self): + ''' + Assemble arguments and store as attributes of the application: + .starter_models: dict of model names to install from INITIAL_CONFIGURE.yaml + True => Install + False => Remove + .scan_directory: Path to a directory of models to scan and import + .autoscan_on_startup: True if invokeai should scan and import at startup time + .import_model_paths: list of URLs, repo_ids and file paths to import + .convert_to_diffusers: if True, convert legacy checkpoints into diffusers + ''' + # starter models to install/remove + model_names = list(self.initial_models.keys()) + starter_models = dict(map(lambda x: (model_names[x], True), self.models_selected.value)) + if hasattr(self,'previously_installed_models'): + unchecked = [ + self.previously_installed_models.values[x] + for x in range(0,len(self.previously_installed_models.values)) + if x not in self.previously_installed_models.value + ] + starter_models.update( + map(lambda x: (x, False), unchecked) + ) + self.parentApp.starter_models=starter_models + + # load directory and whether to scan on startup + self.parentApp.scan_directory = self.autoload_directory.value + self.parentApp.autoscan_on_startup = self.autoscan_on_startup.value + + # URLs and the like + self.parentApp.import_model_paths = self.import_model_paths.value.split() + self.parentApp.convert_to_diffusers = self.convert_models.value != 0 + +class Log(object): + def __init__(self, writable): + self.writable = writable + + def __enter__(self): + self._stdout = sys.stdout + sys.stdout = self.writable + return self + def __exit__(self, *args): + sys.stdout = self._stdout + +class outputForm(npyscreen.ActionForm): + def create(self): + self.buffer = self.add_widget( + npyscreen.BufferPager, + editable=False, + ) + + def write(self,string): + if string != '\n': + self.buffer.buffer([string]) + + def beforeEditing(self): + myapplication = self.parentApp + with Log(self): + print(f'DEBUG: these models will be removed: {[x for x in myapplication.starter_models if not myapplication.starter_models[x]]}') + print(f'DEBUG: these models will be installed: {[x for x in myapplication.starter_models if myapplication.starter_models[x]]}') + print(f'DEBUG: this directory will be scanned: {myapplication.scan_directory}') + print(f'DEBUG: scan at startup time? {myapplication.autoscan_on_startup}') + print(f'DEBUG: these things will be downloaded: {myapplication.import_model_paths}') + print(f'DEBUG: convert to diffusers? {myapplication.convert_to_diffusers}') + + def on_ok(self): + self.buffer.buffer(['goodbye!']) + self.parentApp.setNextForm(None) self.editing = False class AddModelApplication(npyscreen.NPSAppManaged): @@ -242,316 +283,12 @@ class AddModelApplication(npyscreen.NPSAppManaged): addModelsForm, name="Add/Remove Models", ) - - -# --------------------------------------------- -def recommended_datasets() -> dict: - datasets = dict() - for ds in Datasets.keys(): - if Datasets[ds].get("recommended", False): - datasets[ds] = True - return datasets - - -# --------------------------------------------- -def default_dataset() -> dict: - datasets = dict() - for ds in Datasets.keys(): - if Datasets[ds].get("default", False): - datasets[ds] = True - return datasets - - -# --------------------------------------------- -def all_datasets() -> dict: - datasets = dict() - for ds in Datasets.keys(): - datasets[ds] = True - return datasets - - -# --------------------------------------------- -# look for legacy model.ckpt in models directory and offer to -# normalize its name -def migrate_models_ckpt(): - model_path = os.path.join(Globals.root, Model_dir, Weights_dir) - if not os.path.exists(os.path.join(model_path, "model.ckpt")): - return - new_name = Datasets["stable-diffusion-1.4"]["file"] - print('You seem to have the Stable Diffusion v4.1 "model.ckpt" already installed.') - rename = yes_or_no(f'Ok to rename it to "{new_name}" for future reference?') - if rename: - print(f"model.ckpt => {new_name}") - os.replace( - os.path.join(model_path, "model.ckpt"), os.path.join(model_path, new_name) + self.output = self.addForm( + 'MONITOR_OUTPUT', + outputForm, + name='Model Install Output' ) - -# --------------------------------------------- -def download_weight_datasets( - models: dict, access_token: str, precision: str = "float32" -): - migrate_models_ckpt() - successful = dict() - for mod in models.keys(): - print(f"Downloading {mod}:") - successful[mod] = _download_repo_or_file( - Datasets[mod], access_token, precision=precision - ) - return successful - - -def _download_repo_or_file( - mconfig: DictConfig, access_token: str, precision: str = "float32" -) -> Path: - path = None - if mconfig["format"] == "ckpt": - path = _download_ckpt_weights(mconfig, access_token) - else: - path = _download_diffusion_weights(mconfig, access_token, precision=precision) - if "vae" in mconfig and "repo_id" in mconfig["vae"]: - _download_diffusion_weights( - mconfig["vae"], access_token, precision=precision - ) - 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(Globals.root, Model_dir, Weights_dir) - return hf_download_with_resume( - repo_id=repo_id, - model_dir=cache_dir, - model_name=filename, - access_token=access_token, - ) - - -# --------------------------------------------- -def download_from_hf( - model_class: object, model_name: str, cache_subdir: Path = Path("hub"), **kwargs -): - print("", file=sys.stderr) # to prevent tqdm from overwriting - path = global_cache_dir(cache_subdir) - model = model_class.from_pretrained( - model_name, - cache_dir=path, - resume_download=True, - **kwargs, - ) - model_name = "--".join(("models", *model_name.split("/"))) - return path / model_name if model else None - - -def _download_diffusion_weights( - mconfig: DictConfig, access_token: str, precision: str = "float32" -): - repo_id = mconfig["repo_id"] - model_class = ( - StableDiffusionGeneratorPipeline - if mconfig.get("format", None) == "diffusers" - else AutoencoderKL - ) - extra_arg_list = [{"revision": "fp16"}, {}] if precision == "float16" else [{}] - path = None - for extra_args in extra_arg_list: - try: - path = download_from_hf( - model_class, - repo_id, - cache_subdir="diffusers", - safety_checker=None, - **extra_args, - ) - except OSError as e: - if str(e).startswith("fp16 is not a valid"): - pass - else: - print(f"An unexpected error occurred while downloading the model: {e})") - if path: - break - return path - - -# --------------------------------------------- -def hf_download_with_resume( - repo_id: str, model_dir: str, model_name: str, access_token: str = None -) -> Path: - model_dest = Path(os.path.join(model_dir, model_name)) - os.makedirs(model_dir, exist_ok=True) - - url = hf_hub_url(repo_id, model_name) - - header = {"Authorization": f"Bearer {access_token}"} if access_token else {} - open_mode = "wb" - exist_size = 0 - - if os.path.exists(model_dest): - exist_size = os.path.getsize(model_dest) - header["Range"] = f"bytes={exist_size}-" - open_mode = "ab" - - resp = requests.get(url, headers=header, stream=True) - total = int(resp.headers.get("content-length", 0)) - - if ( - resp.status_code == 416 - ): # "range not satisfiable", which means nothing to return - print(f"* {model_name}: complete file found. Skipping.") - return model_dest - elif resp.status_code != 200: - print(f"** An error occurred during downloading {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, - total=total + exist_size, - unit="iB", - unit_scale=True, - unit_divisor=1000, - ) as bar: - for data in resp.iter_content(chunk_size=1024): - size = file.write(data) - bar.update(size) - except Exception as e: - print(f"An error occurred while downloading {model_name}: {str(e)}") - return None - return model_dest - - -# --------------------------------------------- -def update_config_file(successfully_downloaded: dict, opt: dict): - config_file = ( - Path(opt.config_file) if opt.config_file is not None else Default_config_file - ) - - # In some cases (incomplete setup, etc), the default configs directory might be missing. - # Create it if it doesn't exist. - # this check is ignored if opt.config_file is specified - user is assumed to know what they - # are doing if they are passing a custom config file from elsewhere. - if config_file is Default_config_file and not config_file.parent.exists(): - configs_src = Dataset_path.parent - configs_dest = Default_config_file.parent - shutil.copytree(configs_src, configs_dest, dirs_exist_ok=True) - - yaml = new_config_file_contents(successfully_downloaded, config_file, opt) - - try: - backup = None - if os.path.exists(config_file): - print( - 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 - if sys.platform == "win32" and backup.is_file(): - backup.unlink() - config_file.rename(backup) - - with TemporaryFile() as tmp: - tmp.write(Config_preamble.encode()) - tmp.write(yaml.encode()) - - with open(str(config_file.expanduser().resolve()), "wb") as new_config: - tmp.seek(0) - new_config.write(tmp.read()) - - except Exception as e: - print(f"**Error creating config file {config_file}: {str(e)} **") - if backup is not None: - print("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}") - - -# --------------------------------------------- -def new_config_file_contents( - successfully_downloaded: dict, config_file: Path, opt: dict -) -> str: - if config_file.exists(): - conf = OmegaConf.load(str(config_file.expanduser().resolve())) - else: - conf = OmegaConf.create() - - default_selected = None - for model in successfully_downloaded: - # a bit hacky - what we are doing here is seeing whether a checkpoint - # version of the model was previously defined, and whether the current - # model is a diffusers (indicated with a path) - if conf.get(model) and Path(successfully_downloaded[model]).is_dir(): - offer_to_delete_weights(model, conf[model], opt.yes_to_all) - - stanza = {} - mod = Datasets[model] - stanza["description"] = mod["description"] - stanza["repo_id"] = mod["repo_id"] - stanza["format"] = mod["format"] - # diffusers don't need width and height (probably .ckpt doesn't either) - # so we no longer require these in INITIAL_MODELS.yaml - if "width" in mod: - stanza["width"] = mod["width"] - if "height" in mod: - stanza["height"] = mod["height"] - if "file" in mod: - stanza["weights"] = os.path.relpath( - successfully_downloaded[model], start=Globals.root - ) - stanza["config"] = os.path.normpath(os.path.join(SD_Configs, mod["config"])) - if "vae" in mod: - if "file" in mod["vae"]: - stanza["vae"] = os.path.normpath( - os.path.join(Model_dir, Weights_dir, mod["vae"]["file"]) - ) - else: - stanza["vae"] = mod["vae"] - if mod.get("default", False): - stanza["default"] = True - default_selected = True - - conf[model] = stanza - - # if no default model was chosen, then we select the first - # one in the list - if not default_selected: - conf[list(successfully_downloaded.keys())[0]]["default"] = True - - return OmegaConf.to_yaml(conf) - - -# --------------------------------------------- -def offer_to_delete_weights(model_name: str, conf_stanza: dict, yes_to_all: bool): - if not (weights := conf_stanza.get("weights")): - return - if re.match("/VAE/", conf_stanza.get("config")): - return - if yes_to_all or yes_or_no( - f"\n** The checkpoint version of {model_name} is superseded by the diffusers version. Delete the original file {weights}?", - default_yes=False, - ): - weights = Path(weights) - if not weights.is_absolute(): - weights = Path(Globals.root) / weights - try: - weights.unlink() - except OSError as e: - print(str(e)) - - # -------------------------------------------------------- def select_and_download_models(opt: Namespace): if opt.default_only: @@ -559,7 +296,14 @@ def select_and_download_models(opt: Namespace): else: myapplication = AddModelApplication() myapplication.run() - models_to_download = dict(map(lambda x: (x, True), myapplication.selected_models)) if myapplication.selected_models else None + if not myapplication.user_cancelled: + print(f'DEBUG: these models will be removed: {[x for x in myapplication.starter_models if not myapplication.starter_models[x]]}') + print(f'DEBUG: these models will be installed: {[x for x in myapplication.starter_models if myapplication.starter_models[x]]}') + print(f'DEBUG: this directory will be scanned: {myapplication.scan_directory}') + print(f'DEBUG: scan at startup time? {myapplication.autoscan_on_startup}') + print(f'DEBUG: these things will be downloaded: {myapplication.import_model_paths}') + print(f'DEBUG: convert to diffusers? {myapplication.convert_to_diffusers}') + sys.exit(0) if not models_to_download: print( @@ -649,6 +393,7 @@ def main(): ) else: print(f"** A layout error has occurred: {str(e)}") + traceback.print_exc() sys.exit(-1) # ------------------------------------- diff --git a/ldm/invoke/config/model_install_util.py b/ldm/invoke/config/model_install_util.py new file mode 100644 index 0000000000..1182815525 --- /dev/null +++ b/ldm/invoke/config/model_install_util.py @@ -0,0 +1,378 @@ +''' +Utility (backend) functions used by model_install.py +''' +import argparse +import os +import re +import shutil +import sys +import traceback +import warnings +from argparse import Namespace +from math import ceil +from pathlib import Path +from tempfile import TemporaryFile + +import npyscreen +import requests +from diffusers import AutoencoderKL +from huggingface_hub import hf_hub_url +from omegaconf import OmegaConf +from omegaconf.dictconfig import DictConfig +from tqdm import tqdm + +import invokeai.configs as configs +from ldm.invoke.devices import choose_precision, choose_torch_device +from ldm.invoke.generator.diffusers_pipeline import StableDiffusionGeneratorPipeline +from ldm.invoke.globals import Globals, global_cache_dir, global_config_dir +from ldm.invoke.config.widgets import MultiSelectColumns + +warnings.filterwarnings("ignore") +import torch + +# --------------------------globals----------------------- +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 = Path(global_config_dir()) / "models.yaml" +SD_Configs = Path(global_config_dir()) / "stable-diffusion" + +Datasets = OmegaConf.load(Dataset_path) + +Config_preamble = """# This file describes the alternative machine learning models +# available to InvokeAI script. +# +# To add a new model, follow the examples below. Each +# model requires a model config file, a weights file, +# and the width and height of the images it +# was trained on. +""" + +# ------------------------------------- +def yes_or_no(prompt: str, default_yes=True): + default = "y" if default_yes else "n" + response = input(f"{prompt} [{default}] ") or default + if default_yes: + return response[0] not in ("n", "N") + 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 Globals.root + +# --------------------------------------------- +def recommended_datasets() -> dict: + datasets = dict() + for ds in Datasets.keys(): + if Datasets[ds].get("recommended", False): + datasets[ds] = True + return datasets + + +# --------------------------------------------- +def default_dataset() -> dict: + datasets = dict() + for ds in Datasets.keys(): + if Datasets[ds].get("default", False): + datasets[ds] = True + return datasets + + +# --------------------------------------------- +def all_datasets() -> dict: + datasets = dict() + for ds in Datasets.keys(): + datasets[ds] = True + return datasets + + +# --------------------------------------------- +# look for legacy model.ckpt in models directory and offer to +# normalize its name +def migrate_models_ckpt(): + model_path = os.path.join(Globals.root, Model_dir, Weights_dir) + if not os.path.exists(os.path.join(model_path, "model.ckpt")): + return + new_name = Datasets["stable-diffusion-1.4"]["file"] + print('You seem to have the Stable Diffusion v4.1 "model.ckpt" already installed.') + rename = yes_or_no(f'Ok to rename it to "{new_name}" for future reference?') + if rename: + print(f"model.ckpt => {new_name}") + os.replace( + os.path.join(model_path, "model.ckpt"), os.path.join(model_path, new_name) + ) + + +# --------------------------------------------- +def download_weight_datasets( + models: dict, access_token: str, precision: str = "float32" +): + migrate_models_ckpt() + successful = dict() + for mod in models.keys(): + print(f"Downloading {mod}:") + successful[mod] = _download_repo_or_file( + Datasets[mod], access_token, precision=precision + ) + return successful + + +def _download_repo_or_file( + mconfig: DictConfig, access_token: str, precision: str = "float32" +) -> Path: + path = None + if mconfig["format"] == "ckpt": + path = _download_ckpt_weights(mconfig, access_token) + else: + path = _download_diffusion_weights(mconfig, access_token, precision=precision) + if "vae" in mconfig and "repo_id" in mconfig["vae"]: + _download_diffusion_weights( + mconfig["vae"], access_token, precision=precision + ) + 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(Globals.root, Model_dir, Weights_dir) + return hf_download_with_resume( + repo_id=repo_id, + model_dir=cache_dir, + model_name=filename, + access_token=access_token, + ) + + +# --------------------------------------------- +def download_from_hf( + model_class: object, model_name: str, cache_subdir: Path = Path("hub"), **kwargs +): + print("", file=sys.stderr) # to prevent tqdm from overwriting + path = global_cache_dir(cache_subdir) + model = model_class.from_pretrained( + model_name, + cache_dir=path, + resume_download=True, + **kwargs, + ) + model_name = "--".join(("models", *model_name.split("/"))) + return path / model_name if model else None + + +def _download_diffusion_weights( + mconfig: DictConfig, access_token: str, precision: str = "float32" +): + repo_id = mconfig["repo_id"] + model_class = ( + StableDiffusionGeneratorPipeline + if mconfig.get("format", None) == "diffusers" + else AutoencoderKL + ) + extra_arg_list = [{"revision": "fp16"}, {}] if precision == "float16" else [{}] + path = None + for extra_args in extra_arg_list: + try: + path = download_from_hf( + model_class, + repo_id, + cache_subdir="diffusers", + safety_checker=None, + **extra_args, + ) + except OSError as e: + if str(e).startswith("fp16 is not a valid"): + pass + else: + print(f"An unexpected error occurred while downloading the model: {e})") + if path: + break + return path + + +# --------------------------------------------- +def hf_download_with_resume( + repo_id: str, model_dir: str, model_name: str, access_token: str = None +) -> Path: + model_dest = Path(os.path.join(model_dir, model_name)) + os.makedirs(model_dir, exist_ok=True) + + url = hf_hub_url(repo_id, model_name) + + header = {"Authorization": f"Bearer {access_token}"} if access_token else {} + open_mode = "wb" + exist_size = 0 + + if os.path.exists(model_dest): + exist_size = os.path.getsize(model_dest) + header["Range"] = f"bytes={exist_size}-" + open_mode = "ab" + + resp = requests.get(url, headers=header, stream=True) + total = int(resp.headers.get("content-length", 0)) + + if ( + resp.status_code == 416 + ): # "range not satisfiable", which means nothing to return + print(f"* {model_name}: complete file found. Skipping.") + return model_dest + elif resp.status_code != 200: + print(f"** An error occurred during downloading {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, + total=total + exist_size, + unit="iB", + unit_scale=True, + unit_divisor=1000, + ) as bar: + for data in resp.iter_content(chunk_size=1024): + size = file.write(data) + bar.update(size) + except Exception as e: + print(f"An error occurred while downloading {model_name}: {str(e)}") + return None + return model_dest + + +# --------------------------------------------- +def update_config_file(successfully_downloaded: dict, opt: dict): + config_file = ( + Path(opt.config_file) if opt.config_file is not None else Default_config_file + ) + + # In some cases (incomplete setup, etc), the default configs directory might be missing. + # Create it if it doesn't exist. + # this check is ignored if opt.config_file is specified - user is assumed to know what they + # are doing if they are passing a custom config file from elsewhere. + if config_file is Default_config_file and not config_file.parent.exists(): + configs_src = Dataset_path.parent + configs_dest = Default_config_file.parent + shutil.copytree(configs_src, configs_dest, dirs_exist_ok=True) + + yaml = new_config_file_contents(successfully_downloaded, config_file, opt) + + try: + backup = None + if os.path.exists(config_file): + print( + 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 + if sys.platform == "win32" and backup.is_file(): + backup.unlink() + config_file.rename(backup) + + with TemporaryFile() as tmp: + tmp.write(Config_preamble.encode()) + tmp.write(yaml.encode()) + + with open(str(config_file.expanduser().resolve()), "wb") as new_config: + tmp.seek(0) + new_config.write(tmp.read()) + + except Exception as e: + print(f"**Error creating config file {config_file}: {str(e)} **") + if backup is not None: + print("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}") + + +# --------------------------------------------- +def new_config_file_contents( + successfully_downloaded: dict, config_file: Path, opt: dict +) -> str: + if config_file.exists(): + conf = OmegaConf.load(str(config_file.expanduser().resolve())) + else: + conf = OmegaConf.create() + + default_selected = None + for model in successfully_downloaded: + # a bit hacky - what we are doing here is seeing whether a checkpoint + # version of the model was previously defined, and whether the current + # model is a diffusers (indicated with a path) + if conf.get(model) and Path(successfully_downloaded[model]).is_dir(): + offer_to_delete_weights(model, conf[model], opt.yes_to_all) + + stanza = {} + mod = Datasets[model] + stanza["description"] = mod["description"] + stanza["repo_id"] = mod["repo_id"] + stanza["format"] = mod["format"] + # diffusers don't need width and height (probably .ckpt doesn't either) + # so we no longer require these in INITIAL_MODELS.yaml + if "width" in mod: + stanza["width"] = mod["width"] + if "height" in mod: + stanza["height"] = mod["height"] + if "file" in mod: + stanza["weights"] = os.path.relpath( + successfully_downloaded[model], start=Globals.root + ) + stanza["config"] = os.path.normpath(os.path.join(SD_Configs, mod["config"])) + if "vae" in mod: + if "file" in mod["vae"]: + stanza["vae"] = os.path.normpath( + os.path.join(Model_dir, Weights_dir, mod["vae"]["file"]) + ) + else: + stanza["vae"] = mod["vae"] + if mod.get("default", False): + stanza["default"] = True + default_selected = True + + conf[model] = stanza + + # if no default model was chosen, then we select the first + # one in the list + if not default_selected: + conf[list(successfully_downloaded.keys())[0]]["default"] = True + + return OmegaConf.to_yaml(conf) + + +# --------------------------------------------- +def offer_to_delete_weights(model_name: str, conf_stanza: dict, yes_to_all: bool): + if not (weights := conf_stanza.get("weights")): + return + if re.match("/VAE/", conf_stanza.get("config")): + return + if yes_to_all or yes_or_no( + f"\n** The checkpoint version of {model_name} is superseded by the diffusers version. Delete the original file {weights}?", + default_yes=False, + ): + weights = Path(weights) + if not weights.is_absolute(): + weights = Path(Globals.root) / weights + try: + weights.unlink() + except OSError as e: + print(str(e)) + diff --git a/ldm/invoke/config/widgets.py b/ldm/invoke/config/widgets.py index 59fc305ada..b090f83e72 100644 --- a/ldm/invoke/config/widgets.py +++ b/ldm/invoke/config/widgets.py @@ -68,3 +68,30 @@ class MultiSelectColumns(npyscreen.MultiSelect): def h_cursor_line_right(self,ch): super().h_cursor_line_down(ch) + +class TextBox(npyscreen.MultiLineEdit): + def update(self, clear=True): + if clear: self.clear() + + HEIGHT = self.height + WIDTH = self.width + # draw box. + self.parent.curses_pad.hline(self.rely, self.relx, curses.ACS_HLINE, WIDTH) + self.parent.curses_pad.hline(self.rely + HEIGHT, self.relx, curses.ACS_HLINE, WIDTH) + self.parent.curses_pad.vline(self.rely, self.relx, curses.ACS_VLINE, self.height) + self.parent.curses_pad.vline(self.rely, self.relx+WIDTH, curses.ACS_VLINE, HEIGHT) + + # draw corners + self.parent.curses_pad.addch(self.rely, self.relx, curses.ACS_ULCORNER, ) + self.parent.curses_pad.addch(self.rely, self.relx+WIDTH, curses.ACS_URCORNER, ) + self.parent.curses_pad.addch(self.rely+HEIGHT, self.relx, curses.ACS_LLCORNER, ) + self.parent.curses_pad.addch(self.rely+HEIGHT, self.relx+WIDTH, curses.ACS_LRCORNER, ) + + # fool our superclass into thinking drawing area is smaller - this is really hacky but it seems to work + (relx,rely,height,width) = (self.relx, self.rely, self.height, self.width) + self.relx += 1 + self.rely += 1 + self.height -= 1 + self.width -= 1 + super().update(clear=False) + (self.relx,self.rely,self.height,self.width) = (relx, rely, height, width)