mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
3795b40f63
Enhancements: 1. Directory-based imports will not attempt to import components of diffusers models. 2. Diffuser directory imports now supported 3. Files that end with .ckpt that are not Stable Diffusion models (such as VAEs) are skipped during import. Bugs identified in Psychedelicious's review: 1. The invokeai-configure form now tracks the current contents of `invokeai.init` correctly. 2. The autoencoders are no longer treated like installable models, but instead are mandatory support models. They will no longer appear in `models.yaml` Bugs identified in Damian's review: 1. If invokeai-model-install is started before the root directory is initialized, it will call invokeai-configure to fix the matter. 2. Fix bug that was causing empty `models.yaml` under certain conditions. 3. Made import textbox smaller 4. Hide the "convert to diffusers" options if nothing to import.
441 lines
17 KiB
Python
441 lines
17 KiB
Python
#!/usr/bin/env python
|
|
# Copyright (c) 2022 Lincoln D. Stein (https://github.com/lstein)
|
|
# 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.
|
|
|
|
'''
|
|
This is the npyscreen frontend to the model installation application.
|
|
The work is actually done in backend code in model_install_backend.py.
|
|
'''
|
|
|
|
import argparse
|
|
import curses
|
|
import os
|
|
import sys
|
|
import traceback
|
|
from argparse import Namespace
|
|
from typing import List
|
|
|
|
import npyscreen
|
|
import torch
|
|
from pathlib import Path
|
|
from npyscreen import widget
|
|
from omegaconf import OmegaConf
|
|
|
|
from ..devices import choose_precision, choose_torch_device
|
|
from ..globals import Globals, global_config_dir
|
|
from .widgets import MultiSelectColumns, TextBox
|
|
from .model_install_backend import (Dataset_path, default_config_file,
|
|
install_requested_models,
|
|
default_dataset, recommended_datasets, get_root
|
|
)
|
|
|
|
class addModelsForm(npyscreen.FormMultiPage):
|
|
def __init__(self, parentApp, name, multipage=False, *args, **keywords):
|
|
self.multipage=multipage
|
|
self.initial_models = OmegaConf.load(Dataset_path)
|
|
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.installed_models=dict()
|
|
super().__init__(parentApp=parentApp, name=name, *args, **keywords)
|
|
|
|
def create(self):
|
|
window_height, window_width = curses.initscr().getmaxyx()
|
|
starter_model_labels = self._get_starter_model_labels()
|
|
recommended_models = [
|
|
x
|
|
for x in self.starter_model_list
|
|
if self.initial_models[x].get("recommended", False)
|
|
]
|
|
self.installed_models = sorted(
|
|
[
|
|
x for x in list(self.initial_models.keys()) if x in self.existing_models
|
|
]
|
|
)
|
|
self.nextrely -= 1
|
|
self.add_widget_intelligent(
|
|
npyscreen.FixedText,
|
|
value='Use ctrl-N and ctrl-P to move to the <N>ext and <P>revious fields,',
|
|
editable=False,
|
|
)
|
|
self.add_widget_intelligent(
|
|
npyscreen.FixedText,
|
|
value='cursor arrows to make a selection, and space to toggle checkboxes.',
|
|
editable=False,
|
|
)
|
|
self.nextrely += 1
|
|
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",
|
|
)
|
|
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
|
|
)
|
|
self.nextrely += 1
|
|
self.add_widget_intelligent(
|
|
npyscreen.TitleFixedText,
|
|
name="== STARTER MODELS (recommended ones selected) ==",
|
|
value="Select from a starter set of Stable Diffusion models from HuggingFace:",
|
|
begin_entry_at=2,
|
|
editable=False,
|
|
color="CONTROL",
|
|
)
|
|
self.nextrely -= 1
|
|
# if user has already installed some initial models, then don't patronize them
|
|
# by showing more recommendations
|
|
show_recommended = not self.existing_models
|
|
self.models_selected = self.add_widget_intelligent(
|
|
npyscreen.MultiSelect,
|
|
name="Install Starter Models",
|
|
values=starter_model_labels,
|
|
value=[
|
|
self.starter_model_list.index(x)
|
|
for x in self.starter_model_list
|
|
if show_recommended and x in recommended_models
|
|
],
|
|
max_height=len(starter_model_labels) + 1,
|
|
relx = 4,
|
|
scroll_exit=True,
|
|
)
|
|
for line in [
|
|
'== IMPORT LOCAL AND REMOTE MODELS ==',
|
|
'Enter URLs, file paths, or HuggingFace diffusers repository IDs separated by spaces.',
|
|
'Use control-V or shift-control-V to paste:'
|
|
]:
|
|
self.add_widget_intelligent(
|
|
npyscreen.TitleText,
|
|
name=line,
|
|
editable=False,
|
|
color="CONTROL",
|
|
)
|
|
self.nextrely -= 1
|
|
self.import_model_paths = self.add_widget_intelligent(
|
|
TextBox,
|
|
max_height=5,
|
|
scroll_exit=True,
|
|
editable=True,
|
|
relx=4
|
|
)
|
|
self.nextrely += 1
|
|
self.show_directory_fields= self.add_widget_intelligent(
|
|
npyscreen.FormControlCheckbox,
|
|
name='Select a directory for models to import',
|
|
value=False,
|
|
)
|
|
self.autoload_directory = self.add_widget_intelligent(
|
|
npyscreen.TitleFilename,
|
|
name='Directory (<tab> autocompletes):',
|
|
select_dir=True,
|
|
must_exist=True,
|
|
use_two_lines=False,
|
|
labelColor='DANGER',
|
|
begin_entry_at=34,
|
|
scroll_exit=True,
|
|
)
|
|
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
|
|
self.convert_models = self.add_widget_intelligent(
|
|
npyscreen.TitleSelectOne,
|
|
name='== CONVERT IMPORTED MODELS INTO DIFFUSERS==',
|
|
values=['Keep original format','Convert to diffusers'],
|
|
value=0,
|
|
begin_entry_at=4,
|
|
max_height=4,
|
|
hidden=True, # will appear when imported models box is edited
|
|
scroll_exit=True,
|
|
)
|
|
self.cancel = self.add_widget_intelligent(
|
|
npyscreen.ButtonPress,
|
|
name='CANCEL',
|
|
rely = -3,
|
|
when_pressed_function=self.on_cancel,
|
|
)
|
|
done_label = 'DONE'
|
|
back_label = 'BACK'
|
|
button_length = len(done_label)
|
|
button_offset = 0
|
|
if self.multipage:
|
|
button_length += len(back_label)+1
|
|
button_offset += len(back_label)+1
|
|
self.back_button = self.add_widget_intelligent(
|
|
npyscreen.ButtonPress,
|
|
name=back_label,
|
|
relx= (window_width-button_length)//2,
|
|
rely= -3,
|
|
when_pressed_function=self.on_back
|
|
)
|
|
self.ok_button = self.add_widget_intelligent(
|
|
npyscreen.ButtonPress,
|
|
name=done_label,
|
|
relx= button_offset + 1 + (window_width-button_length)//2,
|
|
rely= -3,
|
|
when_pressed_function=self.on_ok
|
|
)
|
|
|
|
for i in [self.autoload_directory,self.autoscan_on_startup]:
|
|
self.show_directory_fields.addVisibleWhenSelected(i)
|
|
|
|
self.show_directory_fields.when_value_edited = self._clear_scan_directory
|
|
self.import_model_paths.when_value_edited = self._show_hide_convert
|
|
self.autoload_directory.when_value_edited = self._show_hide_convert
|
|
|
|
def resize(self):
|
|
super().resize()
|
|
self.models_selected.values = self._get_starter_model_labels()
|
|
|
|
def _clear_scan_directory(self):
|
|
if not self.show_directory_fields.value:
|
|
self.autoload_directory.value = ''
|
|
|
|
def _show_hide_convert(self):
|
|
model_paths = self.import_model_paths.value or ''
|
|
autoload_directory = self.autoload_directory.value or ''
|
|
self.convert_models.hidden = len(model_paths)==0 and len(autoload_directory)==0
|
|
|
|
def _get_starter_model_labels(self)->List[str]:
|
|
window_height, window_width = curses.initscr().getmaxyx()
|
|
label_width = 25
|
|
checkbox_width = 4
|
|
spacing_width = 2
|
|
description_width = window_width - label_width - checkbox_width - spacing_width
|
|
im = self.initial_models
|
|
names = self.starter_model_list
|
|
descriptions = [im[x].description [0:description_width-3]+'...'
|
|
if len(im[x].description) > description_width
|
|
else im[x].description
|
|
for x in names]
|
|
return [
|
|
f"%-{label_width}s %s" % (names[x], descriptions[x]) for x in range(0,len(names))
|
|
]
|
|
|
|
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.editing = False
|
|
self.parentApp.user_cancelled = False
|
|
self.marshall_arguments()
|
|
|
|
def on_back(self):
|
|
self.parentApp.switchFormPrevious()
|
|
self.editing = False
|
|
|
|
def on_cancel(self):
|
|
if npyscreen.notify_yes_no(
|
|
'Are you sure you want to cancel?\nYou may re-run this script later using the invoke.sh or invoke.bat command.\n'
|
|
):
|
|
self.parentApp.setNextForm(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
|
|
'''
|
|
# we're using a global here rather than storing the result in the parentapp
|
|
# due to some bug in npyscreen that is causing attributes to be lost
|
|
selections = self.parentApp.user_selections
|
|
|
|
# starter models to install/remove
|
|
starter_models = dict(map(lambda x: (self.starter_model_list[x], True), self.models_selected.value))
|
|
selections.purge_deleted_models=False
|
|
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)
|
|
)
|
|
selections.purge_deleted_models = self.purge_deleted.value
|
|
selections.starter_models=starter_models
|
|
|
|
# load directory and whether to scan on startup
|
|
if self.show_directory_fields.value:
|
|
selections.scan_directory = self.autoload_directory.value
|
|
selections.autoscan_on_startup = self.autoscan_on_startup.value
|
|
else:
|
|
selections.scan_directory = None
|
|
selections.autoscan_on_startup = False
|
|
|
|
# URLs and the like
|
|
selections.import_model_paths = self.import_model_paths.value.split()
|
|
selections.convert_to_diffusers = self.convert_models.value[0] == 1
|
|
|
|
class AddModelApplication(npyscreen.NPSAppManaged):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.user_cancelled = False
|
|
self.user_selections = Namespace(
|
|
starter_models = None,
|
|
purge_deleted_models = False,
|
|
scan_directory = None,
|
|
autoscan_on_startup = None,
|
|
import_model_paths = None,
|
|
convert_to_diffusers = None
|
|
)
|
|
|
|
def onStart(self):
|
|
npyscreen.setTheme(npyscreen.Themes.DefaultTheme)
|
|
self.main_form = self.addForm(
|
|
"MAIN",
|
|
addModelsForm,
|
|
name="Install Stable Diffusion Models"
|
|
)
|
|
|
|
# --------------------------------------------------------
|
|
def process_and_execute(opt: Namespace, selections: Namespace):
|
|
models_to_remove = [x for x in selections.starter_models if not selections.starter_models[x]]
|
|
models_to_install = [x for x in selections.starter_models if selections.starter_models[x]]
|
|
directory_to_scan = selections.scan_directory
|
|
scan_at_startup = selections.autoscan_on_startup
|
|
potential_models_to_install = selections.import_model_paths
|
|
convert_to_diffusers = selections.convert_to_diffusers
|
|
|
|
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,
|
|
convert_to_diffusers = convert_to_diffusers,
|
|
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,
|
|
)
|
|
|
|
# --------------------------------------------------------
|
|
def select_and_download_models(opt: Namespace):
|
|
precision= 'float32' if opt.full_precision else choose_precision(torch.device(choose_torch_device()))
|
|
if opt.default_only:
|
|
install_requested_models(
|
|
install_initial_models = default_dataset(),
|
|
precision = precision,
|
|
)
|
|
elif opt.yes_to_all:
|
|
install_requested_models(
|
|
install_initial_models = recommended_datasets(),
|
|
precision = precision,
|
|
)
|
|
else:
|
|
installApp = AddModelApplication()
|
|
installApp.run()
|
|
|
|
if not installApp.user_cancelled:
|
|
process_and_execute(opt, installApp.user_selections)
|
|
|
|
# -------------------------------------
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="InvokeAI model downloader")
|
|
parser.add_argument(
|
|
"--full-precision",
|
|
dest="full_precision",
|
|
action=argparse.BooleanOptionalAction,
|
|
type=bool,
|
|
default=False,
|
|
help="use 32-bit weights instead of faster 16-bit weights",
|
|
)
|
|
parser.add_argument(
|
|
"--yes",
|
|
"-y",
|
|
dest="yes_to_all",
|
|
action="store_true",
|
|
help='answer "yes" to all prompts',
|
|
)
|
|
parser.add_argument(
|
|
"--default_only",
|
|
action="store_true",
|
|
help="only install the default model",
|
|
)
|
|
parser.add_argument(
|
|
"--config_file",
|
|
"-c",
|
|
dest="config_file",
|
|
type=str,
|
|
default=None,
|
|
help="path to configuration file to create",
|
|
)
|
|
parser.add_argument(
|
|
"--root_dir",
|
|
dest="root",
|
|
type=str,
|
|
default=None,
|
|
help="path to root of install directory",
|
|
)
|
|
opt = parser.parse_args()
|
|
|
|
# setting a global here
|
|
Globals.root = os.path.expanduser(get_root(opt.root) or "")
|
|
|
|
if not global_config_dir().exists():
|
|
print('>> Your InvokeAI root directory is not set up. Calling invokeai-configure.')
|
|
import ldm.invoke.config.invokeai_configure
|
|
ldm.invoke.config.invokeai_configure.main()
|
|
sys.exit(0)
|
|
|
|
try:
|
|
select_and_download_models(opt)
|
|
except AssertionError as e:
|
|
print(str(e))
|
|
sys.exit(-1)
|
|
except KeyboardInterrupt:
|
|
print("\nGoodbye! Come back soon.")
|
|
except (widget.NotEnoughSpaceForWidget, Exception) as e:
|
|
if str(e).startswith("Height of 1 allocated"):
|
|
print(
|
|
"** Insufficient vertical space for the interface. Please make your window taller and try again"
|
|
)
|
|
elif str(e).startswith('addwstr'):
|
|
print(
|
|
'** Insufficient horizontal space for the interface. Please make your window wider and try again.'
|
|
)
|
|
else:
|
|
print(f"** An error has occurred: {str(e)}")
|
|
traceback.print_exc()
|
|
sys.exit(-1)
|
|
|
|
# -------------------------------------
|
|
if __name__ == "__main__":
|
|
main()
|