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.
This commit is contained in:
d8ahazard 2022-11-08 10:03:56 -06:00
parent ac08562854
commit cfcadeae9a
2 changed files with 26 additions and 2 deletions

View File

@ -1,12 +1,12 @@
import os import os
import sys import sys
import traceback import traceback
from importlib.machinery import SourceFileLoader
import git import git
from modules import paths, shared from modules import paths, shared
extensions = [] extensions = []
extensions_dir = os.path.join(paths.script_path, "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) extension = Extension(name=dirname, path=path, enabled=dirname not in shared.opts.disabled_extensions)
extensions.append(extension) 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

View File

@ -15,7 +15,7 @@ import modules.memmon
import modules.sd_models import modules.sd_models
import modules.styles import modules.styles
import modules.devices as devices 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.hypernetworks import hypernetwork
from modules.paths import models_path, script_path, sd_path 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("--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) parser.add_argument("--server-name", type=str, help="Sets hostname of server", default=None)
extensions.preload_extensions(parser)
cmd_opts = parser.parse_args() cmd_opts = parser.parse_args()
restricted_opts = { restricted_opts = {
"samples_filename_pattern", "samples_filename_pattern",
"directories_filename_pattern", "directories_filename_pattern",