Merge pull request #4488 from d8ahazard/ExtensionPreload

Add option to preload extensions

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 caabf078..17132e42 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",