diff --git a/modules/cmd_args.py b/modules/cmd_args.py index d71982b2c12..7a32cd31f31 100644 --- a/modules/cmd_args.py +++ b/modules/cmd_args.py @@ -31,6 +31,7 @@ parser.add_argument("--no-half-vae", action='store_true', help="do not switch the VAE model to 16-bit floats") parser.add_argument("--no-progressbar-hiding", action='store_true', help="do not hide progressbar in gradio UI (we hide it because it slows down ML if you have hardware acceleration in browser)") parser.add_argument("--max-batch-count", type=int, default=16, help="does not do anything") +parser.add_argument("--max-install-thread", type=int, default=1, help="Maximum Thread number for for asynchronously install extensions. 1 = normal install. ⚠ Enabling this feature may cause unintended issues with some extensions."), parser.add_argument("--embeddings-dir", type=normalized_filepath, default=os.path.join(data_path, 'embeddings'), help="embeddings directory for textual inversion (default: embeddings)") parser.add_argument("--textual-inversion-templates-dir", type=normalized_filepath, default=os.path.join(script_path, 'textual_inversion_templates'), help="directory with textual inversion templates") parser.add_argument("--hypernetwork-dir", type=normalized_filepath, default=os.path.join(models_path, 'hypernetworks'), help="hypernetwork directory") diff --git a/modules/launch_utils.py b/modules/launch_utils.py index 20c7dc127a7..fdc9b878149 100644 --- a/modules/launch_utils.py +++ b/modules/launch_utils.py @@ -12,6 +12,7 @@ import shlex from functools import lru_cache +from concurrent.futures import ThreadPoolExecutor, as_completed from modules import cmd_args, errors from modules.paths_internal import script_path, extensions_dir from modules.timer import startup_timer @@ -228,7 +229,10 @@ def version_check(commit): def run_extension_installer(extension_dir): path_installer = os.path.join(extension_dir, "install.py") if not os.path.isfile(path_installer): - return + return False + + dirname = os.path.basename(extension_dir) + logging.debug(f"Installing {dirname}") try: env = os.environ.copy() @@ -240,6 +244,8 @@ def run_extension_installer(extension_dir): except Exception as e: errors.report(str(e)) + return True + def list_extensions(settings_file): settings = {} @@ -267,14 +273,26 @@ def run_extensions_installers(settings_file): return with startup_timer.subcategory("run extensions installers"): + paths = {} for dirname_extension in list_extensions(settings_file): - logging.debug(f"Installing {dirname_extension}") - path = os.path.join(extensions_dir, dirname_extension) if os.path.isdir(path): + paths[dirname_extension] = path + + max_workers = args.max_install_thread + if max_workers == 1: + for dirname_extension, path in paths.items(): run_extension_installer(path) startup_timer.record(dirname_extension) + else: + max_workers = min(max_workers, 4) + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = {executor.submit(run_extension_installer, path): dirname_extension for dirname_extension, path in paths.items()} + for future in as_completed(futures): + dirname_extension = futures[future] + if future.result(): + startup_timer.record(dirname_extension) re_requirement = re.compile(r"\s*([-_a-zA-Z0-9]+)\s*(?:==\s*([-+_.a-zA-Z0-9]+))?\s*")