autoimport from embedding/controlnet/lora folders designated in startup file

This commit is contained in:
Lincoln Stein
2023-06-27 12:30:53 -04:00
parent f15d28d141
commit e8ed0fad6c
7 changed files with 172 additions and 123 deletions

View File

@ -442,6 +442,26 @@ to allow InvokeAI to download restricted styles & subjects from the "Concept Lib
scroll_exit=True,
)
self.nextrely += 1
self.add_widget_intelligent(
npyscreen.FixedText,
value="Directories containing textual inversion, controlnet and LoRA models (<tab> autocompletes, ctrl-N advances):",
editable=False,
color="CONTROL",
)
self.autoimport_dirs = {}
for description, config_name, path in autoimport_paths(old_opts):
self.autoimport_dirs[config_name] = self.add_widget_intelligent(
npyscreen.TitleFilename,
name=description+':',
value=str(path),
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,
name="== LICENSE ==",
@ -505,10 +525,6 @@ https://huggingface.co/spaces/CompVis/stable-diffusion-license
bad_fields.append(
f"The output directory does not seem to be valid. Please check that {str(Path(opt.outdir).parent)} is an existing directory."
)
# if not Path(opt.embedding_dir).parent.exists():
# bad_fields.append(
# f"The embedding directory does not seem to be valid. Please check that {str(Path(opt.embedding_dir).parent)} is an existing directory."
# )
if len(bad_fields) > 0:
message = "The following problems were detected and must be corrected:\n"
for problem in bad_fields:
@ -528,12 +544,15 @@ https://huggingface.co/spaces/CompVis/stable-diffusion-license
"max_loaded_models",
"xformers_enabled",
"always_use_cpu",
# "embedding_dir",
# "lora_dir",
# "controlnet_dir",
]:
setattr(new_opts, attr, getattr(self, attr).value)
for attr in self.autoimport_dirs:
directory = Path(self.autoimport_dirs[attr].value)
if directory.is_relative_to(config.root_path):
directory = directory.relative_to(config.root_path)
setattr(new_opts, attr, directory)
new_opts.hf_token = self.hf_token.value
new_opts.license_acceptance = self.license_acceptance.value
new_opts.precision = PRECISION_CHOICES[self.precision.value[0]]
@ -595,22 +614,32 @@ def default_user_selections(program_opts: Namespace) -> InstallSelections:
else [models[x].path or models[x].repo_id for x in installer.recommended_models()]
if program_opts.yes_to_all
else list(),
scan_directory=None,
autoscan_on_startup=None,
# scan_directory=None,
# autoscan_on_startup=None,
)
# -------------------------------------
def autoimport_paths(config: InvokeAIAppConfig):
return [
('Checkpoints & diffusers models', 'autoimport_dir', config.root_path / config.autoimport_dir),
('LoRA/LyCORIS models', 'lora_dir', config.root_path / config.lora_dir),
('Controlnet models', 'controlnet_dir', config.root_path / config.controlnet_dir),
('Textual Inversion Embeddings', 'embedding_dir', config.root_path / config.embedding_dir),
]
# -------------------------------------
def initialize_rootdir(root: Path, yes_to_all: bool = False):
logger.info("** INITIALIZING INVOKEAI RUNTIME DIRECTORY **")
for name in (
"models",
"databases",
"autoimport",
"text-inversion-output",
"text-inversion-training-data",
"configs"
):
os.makedirs(os.path.join(root, name), exist_ok=True)
for model_type in ModelType:
Path(root, 'autoimport', model_type.value).mkdir(parents=True, exist_ok=True)
configs_src = Path(configs.__path__[0])
configs_dest = root / "configs"
@ -618,9 +647,8 @@ def initialize_rootdir(root: Path, yes_to_all: bool = False):
shutil.copytree(configs_src, configs_dest, dirs_exist_ok=True)
dest = root / 'models'
for model_base in [BaseModelType.StableDiffusion1,BaseModelType.StableDiffusion2]:
for model_type in [ModelType.Main, ModelType.Vae, ModelType.Lora,
ModelType.ControlNet,ModelType.TextualInversion]:
for model_base in BaseModelType:
for model_type in ModelType:
path = dest / model_base.value / model_type.value
path.mkdir(parents=True, exist_ok=True)
path = dest / 'core'
@ -632,9 +660,7 @@ def initialize_rootdir(root: Path, yes_to_all: bool = False):
}
)
)
# with open(root / 'invokeai.yaml','w') as f:
# f.write('#empty invokeai.yaml initialization file')
# -------------------------------------
def run_console_ui(
program_opts: Namespace, initfile: Path = None

View File

@ -70,8 +70,8 @@ class ModelInstallList:
class InstallSelections():
install_models: List[str]= field(default_factory=list)
remove_models: List[str]=field(default_factory=list)
scan_directory: Path = None
autoscan_on_startup: bool=False
# scan_directory: Path = None
# autoscan_on_startup: bool=False
@dataclass
class ModelLoadInfo():
@ -155,8 +155,8 @@ class ModelInstall(object):
def install(self, selections: InstallSelections):
job = 1
jobs = len(selections.remove_models) + len(selections.install_models)
if selections.scan_directory:
jobs += 1
# if selections.scan_directory:
# jobs += 1
# remove requested models
for key in selections.remove_models:
@ -171,18 +171,8 @@ class ModelInstall(object):
self.heuristic_install(path)
job += 1
# import from the scan directory, if any
if path := selections.scan_directory:
logger.info(f'Scanning and importing models from directory {path} [{job}/{jobs}]')
self.heuristic_install(path)
self.mgr.commit()
if selections.autoscan_on_startup and Path(selections.scan_directory).is_dir():
update_autoimport_dir(selections.scan_directory)
else:
update_autoimport_dir(None)
def heuristic_install(self,
model_path_id_or_url: Union[str,Path],
models_installed: Set[Path]=None)->Set[Path]:
@ -237,7 +227,7 @@ class ModelInstall(object):
self.mgr.add_model(model_name = model_name,
base_model = info.base_type,
model_type = info.model_type,
model_attributes = attributes
model_attributes = attributes,
)
except Exception as e:
logger.warning(f'{str(e)} Skipping registration.')
@ -309,11 +299,11 @@ class ModelInstall(object):
return location.stem
def _make_attributes(self, path: Path, info: ModelProbeInfo)->dict:
# convoluted way to retrieve the description from datasets
description = f'{info.base_type.value} {info.model_type.value} model'
model_name = path.name if path.is_dir() else path.stem
description = f'{info.base_type.value} {info.model_type.value} model {model_name}'
if key := self.reverse_paths.get(self.current_id):
if key in self.datasets:
description = self.datasets[key]['description']
description = self.datasets[key].get('description') or description
rel_path = self.relative_to_root(path)
@ -395,19 +385,6 @@ class ModelInstall(object):
'''
return {v.get('path') or v.get('repo_id') : k for k, v in datasets.items()}
def update_autoimport_dir(autodir: Path):
'''
Update the "autoimport_dir" option in invokeai.yaml
'''
invokeai_config_path = config.init_file_path
conf = OmegaConf.load(invokeai_config_path)
conf.InvokeAI.Paths.autoimport_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)
# -------------------------------------
def yes_or_no(prompt: str, default_yes=True):
default = "y" if default_yes else "n"