mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
27a2e27c3a
1. Fixed display crash when the number of installed models is less than the number of desired columns to display them. 2. Added --ckpt_convert option to init file.
442 lines
17 KiB
Python
442 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()
|
|
cols = 4 if window_width > 240 else 3 if window_width>160 else 2 if window_width>80 else 1
|
|
return min(cols,len(self.installed_models))
|
|
|
|
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()
|