From cfcadeae9a61e1aff32960864f90299412c86d5c Mon Sep 17 00:00:00 2001
From: d8ahazard <d8ahazard@gmail.com>
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",