Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

add controlnet model downloading

Browse Source
This commit is contained in:
Lincoln Stein 2023-05-30 13:49:43 -04:00
parent c9ee42450e
commit 1632ac6b9f
3 changed files with 206 additions and 127 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
invokeai/app/services/config.py
View File

@ -359,7 +359,7 @@ setting environment variables INVOKEAI_<setting>.
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_<setting>.
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:
'''

140
invokeai/backend/config/model_install_backend.py
View File

@ -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,

184
invokeai/frontend/install/model_install.py
View File

@ -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,
)
# --------------------------------------------------------