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') 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') 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') 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') 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') 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') 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: def _resolve(self,partial_path:Path)->Path:
return (self.root_path / partial_path).resolve() return (self.root_path / partial_path).resolve()
@property
def init_file_path(self)->Path:
'''
Path to invokeai.yaml
'''
return self._resolve(INIT_FILE)
@property @property
def output_path(self)->Path: def output_path(self)->Path:
''' '''

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

@ -8,11 +8,11 @@ import sys
import warnings import warnings
from pathlib import Path from pathlib import Path
from tempfile import TemporaryFile from tempfile import TemporaryFile
from typing import List from typing import List, Dict
import requests import requests
from diffusers import AutoencoderKL 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 import OmegaConf
from omegaconf.dictconfig import DictConfig from omegaconf.dictconfig import DictConfig
from tqdm import tqdm from tqdm import tqdm
@ -49,7 +49,6 @@ Config_preamble = """
def default_config_file(): def default_config_file():
print(config.root_dir)
return config.model_conf_path return config.model_conf_path
def sd_configs(): def sd_configs():
@ -62,23 +61,35 @@ def initial_models():
return (Datasets := OmegaConf.load(Dataset_path)['diffusers']) return (Datasets := OmegaConf.load(Dataset_path)['diffusers'])
def install_requested_models( def install_requested_models(
install_initial_models: List[str] = None, install_initial_models: List[str] = None,
remove_models: List[str] = None, remove_models: List[str] = None,
scan_directory: Path = None, install_cn_models: List[str] = None,
external_models: List[str] = None, remove_cn_models: List[str] = None,
scan_at_startup: bool = False, cn_model_map: Dict[str,str] = None,
precision: str = "float16", scan_directory: Path = None,
purge_deleted: bool = False, external_models: List[str] = None,
config_file_path: Path = 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. 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() config_file_path = config_file_path or default_config_file()
if not config_file_path.exists(): if not config_file_path.exists():
open(config_file_path, "w") 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: if remove_models and len(remove_models) > 0:
print("== DELETING UNCHECKED STARTER MODELS ==") print("== DELETING UNCHECKED STARTER MODELS ==")
@ -120,18 +131,20 @@ def install_requested_models(
pass pass
if scan_at_startup and scan_directory.is_dir(): if scan_at_startup and scan_directory.is_dir():
argument = "--autoconvert" update_autoconvert_dir(scan_directory)
print('** The global initfile is no longer supported; rewrite to support new yaml format **')
initfile = Path(config.root_dir, 'invokeai.init') def update_autoconvert_dir(autodir: Path):
replacement = Path(config.root_dir, f"invokeai.init.new") '''
directory = str(scan_directory).replace("\\", "/") Update the "autoconvert_dir" option in invokeai.yaml
with open(initfile, "r") as input: '''
with open(replacement, "w") as output: invokeai_config_path = config.init_file_path
while line := input.readline(): conf = OmegaConf.load(invokeai_config_path)
if not line.startswith(argument): conf.InvokeAI.Paths.autoconvert_dir = str(autodir)
output.writelines([line]) yaml = OmegaConf.to_yaml(conf)
output.writelines([f"{argument} {directory}"]) tmpfile = invokeai_config_path.parent / "new_config.tmp"
os.replace(replacement, initfile) 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( def download_from_hf(
model_class: object, model_name: str, **kwargs model_class: object, model_name: str, **kwargs
@ -273,9 +348,13 @@ def _download_diffusion_weights(
# --------------------------------------------- # ---------------------------------------------
def hf_download_with_resume( 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: ) -> 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) os.makedirs(model_dir, exist_ok=True)
url = hf_hub_url(repo_id, model_name) 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 ): # "range not satisfiable", which means nothing to return
print(f"* {model_name}: complete file found. Skipping.") print(f"* {model_name}: complete file found. Skipping.")
return model_dest return model_dest
elif resp.status_code == 404:
print("** File not found")
return None
elif resp.status_code != 200: 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: elif exist_size > 0:
print(f"* {model_name}: partial file found. Resuming...") print(f"* {model_name}: partial file found. Resuming...")
else: else:
print(f"* {model_name}: Downloading...") print(f"* {model_name}: Downloading...")
try: try:
if total < 2000:
print(f"*** ERROR DOWNLOADING {model_name}: {resp.text}")
return None
with open(model_dest, open_mode) as file, tqdm( with open(model_dest, open_mode) as file, tqdm(
desc=model_name, desc=model_name,
initial=exist_size, initial=exist_size,

182
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 # minimum size for the UI
MIN_COLS = 120 MIN_COLS = 120
MIN_LINES = 45 MIN_LINES = 50
config = get_invokeai_config() config = get_invokeai_config()
@ -53,16 +53,16 @@ class addModelsForm(npyscreen.FormMultiPage):
def __init__(self, parentApp, name, multipage=False, *args, **keywords): def __init__(self, parentApp, name, multipage=False, *args, **keywords):
self.multipage = multipage self.multipage = multipage
self.initial_models = OmegaConf.load(Dataset_path)['diffusers'] self.initial_models = OmegaConf.load(Dataset_path)['diffusers']
self.control_net_models = OmegaConf.load(Dataset_path)['controlnet'] self.control_net_models = OmegaConf.load(Dataset_path)['controlnet']
self.installed_cn_models = self._get_installed_cn_models() self.installed_cn_models = self._get_installed_cn_models()
self._add_additional_cn_models(self.control_net_models,self.installed_cn_models)
try: try:
self.existing_models = OmegaConf.load(default_config_file()) self.existing_models = OmegaConf.load(default_config_file())
except: except:
self.existing_models = dict() 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.starter_model_list = list(self.initial_models.keys())
self.installed_models = dict() self.installed_models = dict()
super().__init__(parentApp=parentApp, name=name, *args, **keywords) super().__init__(parentApp=parentApp, name=name, *args, **keywords)
@ -95,40 +95,6 @@ class addModelsForm(npyscreen.FormMultiPage):
color="CAUTION", color="CAUTION",
) )
self.nextrely += 1 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: if len(self.starter_model_list) > 0:
self.add_widget_intelligent( self.add_widget_intelligent(
CenteredTitleText, CenteredTitleText,
@ -161,35 +127,14 @@ class addModelsForm(npyscreen.FormMultiPage):
relx=4, relx=4,
scroll_exit=True, 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( self.purge_deleted = self.add_widget_intelligent(
npyscreen.Checkbox, npyscreen.Checkbox,
name="Purge unchecked models from disk", name="Purge unchecked diffusers models from disk",
value=False, value=False,
scroll_exit=True, scroll_exit=True,
relx=4, relx=4,
) )
self.nextrely += 1
self.add_widget_intelligent( self.add_widget_intelligent(
CenteredTitleText, CenteredTitleText,
name="== IMPORT LOCAL AND REMOTE MODELS ==", name="== IMPORT LOCAL AND REMOTE MODELS ==",
@ -211,7 +156,7 @@ class addModelsForm(npyscreen.FormMultiPage):
) )
self.nextrely -= 1 self.nextrely -= 1
self.import_model_paths = self.add_widget_intelligent( 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.nextrely += 1
self.show_directory_fields = self.add_widget_intelligent( self.show_directory_fields = self.add_widget_intelligent(
@ -236,6 +181,47 @@ class addModelsForm(npyscreen.FormMultiPage):
relx=4, relx=4,
scroll_exit=True, 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( self.cancel = self.add_widget_intelligent(
npyscreen.ButtonPress, npyscreen.ButtonPress,
name="CANCEL", name="CANCEL",
@ -300,19 +286,21 @@ class addModelsForm(npyscreen.FormMultiPage):
] ]
def _get_installed_cn_models(self)->list[str]: def _get_installed_cn_models(self)->list[str]:
with open('log.txt','w') as file: cn_dir = config.controlnet_path
cn_dir = config.controlnet_path installed_cn_models = set()
file.write(f'cn_dir={cn_dir}\n') for root, dirs, files in os.walk(cn_dir):
installed_cn_models = set() for name in dirs:
for root, dirs, files in os.walk(cn_dir): if Path(root, name, '.download_complete').exists():
for name in dirs: installed_cn_models.add(name.replace('--','/'))
file.write(f'{root}/{name}/config.json\n') return installed_cn_models
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]
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: def _get_columns(self) -> int:
window_width, window_height = get_terminal_size() window_width, window_height = get_terminal_size()
@ -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.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.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 for x in self.cn_models_selected.value
if self.cn_models_selected.values[x] not in self.installed_cn_models 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 for x in self.cn_models_selected.values
if x in self.installed_cn_models if x in self.installed_cn_models
and self.cn_models_selected.values.index(x) not in self.cn_models_selected.value 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 # load directory and whether to scan on startup
if self.show_directory_fields.value: if self.show_directory_fields.value:
@ -406,6 +399,7 @@ class AddModelApplication(npyscreen.NPSAppManaged):
purge_deleted_models=False, purge_deleted_models=False,
install_cn_models = None, install_cn_models = None,
remove_cn_models = None, remove_cn_models = None,
control_net_map = None,
scan_directory=None, scan_directory=None,
autoscan_on_startup=None, autoscan_on_startup=None,
import_model_paths=None, import_model_paths=None,
@ -425,24 +419,24 @@ def process_and_execute(opt: Namespace, selections: Namespace):
directory_to_scan = selections.scan_directory directory_to_scan = selections.scan_directory
scan_at_startup = selections.autoscan_on_startup scan_at_startup = selections.autoscan_on_startup
potential_models_to_install = selections.import_model_paths potential_models_to_install = selections.import_model_paths
print(f'selections.install_cn_models={selections.install_cn_models}')
print('NOT INSTALLING MODELS DURING DEBUGGING') print(f'selections.remove_cn_models={selections.remove_cn_models}')
print('models to install:',models_to_install) print(f'selections.cn_model_map={selections.control_net_map}')
print('models to remove:',models_to_remove) install_requested_models(
print('CN models to install:',selections.install_cn_models) install_initial_models=models_to_install,
print('CN models to remove:',selections.remove_cn_models) remove_models=models_to_remove,
# install_requested_models( install_cn_models=selections.install_cn_models,
# install_initial_models=models_to_install, remove_cn_models=selections.remove_cn_models,
# remove_models=models_to_remove, cn_model_map=selections.control_net_map,
# scan_directory=Path(directory_to_scan) if directory_to_scan else None, scan_directory=Path(directory_to_scan) if directory_to_scan else None,
# external_models=potential_models_to_install, external_models=potential_models_to_install,
# scan_at_startup=scan_at_startup, scan_at_startup=scan_at_startup,
# precision="float32" precision="float32"
# if opt.full_precision if opt.full_precision
# else choose_precision(torch.device(choose_torch_device())), else choose_precision(torch.device(choose_torch_device())),
# purge_deleted=selections.purge_deleted_models, purge_deleted=selections.purge_deleted_models,
# config_file_path=Path(opt.config_file) if opt.config_file else None, config_file_path=Path(opt.config_file) if opt.config_file else None,
# ) )
# -------------------------------------------------------- # --------------------------------------------------------