From cfcadeae9a61e1aff32960864f90299412c86d5c Mon Sep 17 00:00:00 2001 From: d8ahazard Date: Tue, 8 Nov 2022 10:03:56 -0600 Subject: [PATCH] Add option to preload extensions By creating a file called "preload.py" in an extension folder and declaring a preload(parser) method, we can add extra command-line args for an extension. --- modules/extensions.py | 23 ++++++++++++++++++++++- modules/shared.py | 5 ++++- 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/modules/extensions.py b/modules/extensions.py index 8e0977fd..544f3580 100644 --- a/modules/extensions.py +++ b/modules/extensions.py @@ -1,12 +1,12 @@ import os import sys import traceback +from importlib.machinery import SourceFileLoader import git from modules import paths, shared - extensions = [] extensions_dir = os.path.join(paths.script_path, "extensions") @@ -84,3 +84,24 @@ def list_extensions(): extension = Extension(name=dirname, path=path, enabled=dirname not in shared.opts.disabled_extensions) extensions.append(extension) + + +def preload_extensions(parser): + if not os.path.isdir(extensions_dir): + return + + for dirname in sorted(os.listdir(extensions_dir)): + path = os.path.join(extensions_dir, dirname) + if not os.path.isdir(path): + continue + for file in os.listdir(path): + if "preload.py" in file: + full_file = os.path.join(path, file) + print(f"Got preload file: {full_file}") + + try: + ext = SourceFileLoader("preload", full_file).load_module() + parser = ext.preload(parser) + except Exception as e: + print(f"Exception preloading script: {e}") + return parser \ No newline at end of file diff --git a/modules/shared.py b/modules/shared.py index e8bacd3c..222ad4fb 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -15,7 +15,7 @@ import modules.memmon import modules.sd_models import modules.styles import modules.devices as devices -from modules import sd_samplers, sd_models, localization, sd_vae +from modules import sd_samplers, sd_models, localization, sd_vae, extensions from modules.hypernetworks import hypernetwork from modules.paths import models_path, script_path, sd_path @@ -91,7 +91,10 @@ parser.add_argument("--tls-keyfile", type=str, help="Partially enables TLS, requ parser.add_argument("--tls-certfile", type=str, help="Partially enables TLS, requires --tls-keyfile to fully function", default=None) parser.add_argument("--server-name", type=str, help="Sets hostname of server", default=None) +extensions.preload_extensions(parser) + cmd_opts = parser.parse_args() + restricted_opts = { "samples_filename_pattern", "directories_filename_pattern",