diff --git a/launch.py b/launch.py index 5fa11560..8e65676d 100644 --- a/launch.py +++ b/launch.py @@ -105,22 +105,26 @@ def version_check(commit): print("version check failed", e) +def run_extension_installer(extension_dir): + path_installer = os.path.join(extension_dir, "install.py") + if not os.path.isfile(path_installer): + return + + try: + env = os.environ.copy() + env['PYTHONPATH'] = os.path.abspath(".") + + print(run(f'"{python}" "{path_installer}"', errdesc=f"Error running install.py for extension {extension_dir}", custom_env=env)) + except Exception as e: + print(e, file=sys.stderr) + + def run_extensions_installers(): if not os.path.isdir(dir_extensions): return for dirname_extension in os.listdir(dir_extensions): - path_installer = os.path.join(dir_extensions, dirname_extension, "install.py") - if not os.path.isfile(path_installer): - continue - - try: - env = os.environ.copy() - env['PYTHONPATH'] = os.path.abspath(".") - - print(run(f'"{python}" "{path_installer}"', errdesc=f"Error running install.py for extension {dirname_extension}", custom_env=env)) - except Exception as e: - print(e, file=sys.stderr) + run_extension_installer(os.path.join(dir_extensions, dirname_extension)) def prepare_enviroment(): diff --git a/modules/extensions.py b/modules/extensions.py index 544f3580..94ce479a 100644 --- a/modules/extensions.py +++ b/modules/extensions.py @@ -1,7 +1,6 @@ import os import sys import traceback -from importlib.machinery import SourceFileLoader import git @@ -85,23 +84,3 @@ 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/script_loading.py b/modules/script_loading.py new file mode 100644 index 00000000..f93f0951 --- /dev/null +++ b/modules/script_loading.py @@ -0,0 +1,34 @@ +import os +import sys +import traceback +from types import ModuleType + + +def load_module(path): + with open(path, "r", encoding="utf8") as file: + text = file.read() + + compiled = compile(text, path, 'exec') + module = ModuleType(os.path.basename(path)) + exec(compiled, module.__dict__) + + return module + + +def preload_extensions(extensions_dir, parser): + if not os.path.isdir(extensions_dir): + return + + for dirname in sorted(os.listdir(extensions_dir)): + preload_script = os.path.join(extensions_dir, dirname, "preload.py") + if not os.path.isfile(preload_script): + continue + + try: + module = load_module(preload_script) + if hasattr(module, 'preload'): + module.preload(parser) + + except Exception: + print(f"Error running preload() for {preload_script}", file=sys.stderr) + print(traceback.format_exc(), file=sys.stderr) diff --git a/modules/scripts.py b/modules/scripts.py index 22d8908b..986b1914 100644 --- a/modules/scripts.py +++ b/modules/scripts.py @@ -6,7 +6,7 @@ from collections import namedtuple import gradio as gr from modules.processing import StableDiffusionProcessing -from modules import shared, paths, script_callbacks, extensions +from modules import shared, paths, script_callbacks, extensions, script_loading AlwaysVisible = object() @@ -161,13 +161,7 @@ def load_scripts(): sys.path = [scriptfile.basedir] + sys.path current_basedir = scriptfile.basedir - with open(scriptfile.path, "r", encoding="utf8") as file: - text = file.read() - - from types import ModuleType - compiled = compile(text, scriptfile.path, 'exec') - module = ModuleType(scriptfile.filename) - exec(compiled, module.__dict__) + module = script_loading.load_module(scriptfile.path) for key, script_class in module.__dict__.items(): if type(script_class) == type and issubclass(script_class, Script): @@ -328,27 +322,21 @@ class ScriptRunner: def reload_sources(self, cache): for si, script in list(enumerate(self.scripts)): - with open(script.filename, "r", encoding="utf8") as file: - args_from = script.args_from - args_to = script.args_to - filename = script.filename - text = file.read() + args_from = script.args_from + args_to = script.args_to + filename = script.filename - from types import ModuleType + module = cache.get(filename, None) + if module is None: + module = script_loading.load_module(script.filename) + cache[filename] = module - module = cache.get(filename, None) - if module is None: - compiled = compile(text, filename, 'exec') - module = ModuleType(script.filename) - exec(compiled, module.__dict__) - cache[filename] = module - - for key, script_class in module.__dict__.items(): - if type(script_class) == type and issubclass(script_class, Script): - self.scripts[si] = script_class() - self.scripts[si].filename = filename - self.scripts[si].args_from = args_from - self.scripts[si].args_to = args_to + for key, script_class in module.__dict__.items(): + if type(script_class) == type and issubclass(script_class, Script): + self.scripts[si] = script_class() + self.scripts[si].filename = filename + self.scripts[si].args_from = args_from + self.scripts[si].args_to = args_to scripts_txt2img = ScriptRunner() diff --git a/modules/sd_vae.py b/modules/sd_vae.py index d82a7bad..a6310a2e 100644 --- a/modules/sd_vae.py +++ b/modules/sd_vae.py @@ -87,7 +87,19 @@ def refresh_vae_list(vae_path=vae_path, model_path=model_path): return vae_list -def resolve_vae(checkpoint_file, vae_file="auto"): +def get_vae_from_settings(vae_file="auto"): + # else, we load from settings, if not set to be default + if vae_file == "auto" and shared.opts.sd_vae is not None: + # if saved VAE settings isn't recognized, fallback to auto + vae_file = vae_dict.get(shared.opts.sd_vae, "auto") + # if VAE selected but not found, fallback to auto + if vae_file not in default_vae_values and not os.path.isfile(vae_file): + vae_file = "auto" + print("Selected VAE doesn't exist") + return vae_file + + +def resolve_vae(checkpoint_file=None, vae_file="auto"): global first_load, vae_dict, vae_list # if vae_file argument is provided, it takes priority, but not saved @@ -102,14 +114,9 @@ def resolve_vae(checkpoint_file, vae_file="auto"): shared.opts.data['sd_vae'] = get_filename(vae_file) else: print("VAE provided as command line argument doesn't exist") - # else, we load from settings - if vae_file == "auto" and shared.opts.sd_vae is not None: - # if saved VAE settings isn't recognized, fallback to auto - vae_file = vae_dict.get(shared.opts.sd_vae, "auto") - # if VAE selected but not found, fallback to auto - if vae_file not in default_vae_values and not os.path.isfile(vae_file): - vae_file = "auto" - print("Selected VAE doesn't exist") + # fallback to selector in settings, if vae selector not set to act as default fallback + if not shared.opts.sd_vae_as_default: + vae_file = get_vae_from_settings(vae_file) # vae-path cmd arg takes priority for auto if vae_file == "auto" and shared.cmd_opts.vae_path is not None: if os.path.isfile(shared.cmd_opts.vae_path): diff --git a/modules/shared.py b/modules/shared.py index a9daf800..84567c8e 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -3,7 +3,6 @@ import datetime import json import os import sys -from collections import OrderedDict import time import gradio as gr @@ -15,7 +14,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, extensions +from modules import sd_samplers, sd_models, localization, sd_vae, extensions, script_loading from modules.hypernetworks import hypernetwork from modules.paths import models_path, script_path, sd_path @@ -91,7 +90,7 @@ 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) +script_loading.preload_extensions(extensions.extensions_dir, parser) cmd_opts = parser.parse_args() @@ -336,6 +335,7 @@ options_templates.update(options_section(('sd', "Stable Diffusion"), { "sd_model_checkpoint": OptionInfo(None, "Stable Diffusion checkpoint", gr.Dropdown, lambda: {"choices": modules.sd_models.checkpoint_tiles()}, refresh=sd_models.list_models), "sd_checkpoint_cache": OptionInfo(0, "Checkpoints to cache in RAM", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}), "sd_vae": OptionInfo("auto", "SD VAE", gr.Dropdown, lambda: {"choices": sd_vae.vae_list}, refresh=sd_vae.refresh_vae_list), + "sd_vae_as_default": OptionInfo(False, "Ignore selected VAE for stable diffusion checkpoints that have their own .vae.pt next to them"), "sd_hypernetwork": OptionInfo("None", "Hypernetwork", gr.Dropdown, lambda: {"choices": ["None"] + [x for x in hypernetworks.keys()]}, refresh=reload_hypernetworks), "sd_hypernetwork_strength": OptionInfo(1.0, "Hypernetwork strength", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.001}), "inpainting_mask_weight": OptionInfo(1.0, "Inpainting conditioning mask strength", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}), diff --git a/modules/ui_extensions.py b/modules/ui_extensions.py index 02ab9643..6671cb60 100644 --- a/modules/ui_extensions.py +++ b/modules/ui_extensions.py @@ -134,6 +134,9 @@ def install_extension_from_url(dirname, url): os.rename(tmpdir, target_dir) + import launch + launch.run_extension_installer(target_dir) + extensions.list_extensions() return [extension_table(), html.escape(f"Installed into {target_dir}. Use Installed tab to restart.")] finally: diff --git a/requirements.txt b/requirements.txt index 0fba0b23..762db4f3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ +accelerate basicsr diffusers fairscale==0.4.4 diff --git a/requirements_versions.txt b/requirements_versions.txt index f7059f20..662ca684 100644 --- a/requirements_versions.txt +++ b/requirements_versions.txt @@ -1,5 +1,6 @@ transformers==4.19.2 diffusers==0.3.0 +accelerate==0.12.0 basicsr==1.4.2 gfpgan==1.3.8 gradio==3.9 diff --git a/webui-user.bat b/webui-user.bat index e5a257be..0c120111 100644 --- a/webui-user.bat +++ b/webui-user.bat @@ -4,5 +4,6 @@ set PYTHON= set GIT= set VENV_DIR= set COMMANDLINE_ARGS= +set ACCELERATE= call webui.bat diff --git a/webui-user.sh b/webui-user.sh index 30646f5c..16e42759 100644 --- a/webui-user.sh +++ b/webui-user.sh @@ -40,4 +40,7 @@ export COMMANDLINE_ARGS="" #export CODEFORMER_COMMIT_HASH="" #export BLIP_COMMIT_HASH="" +# Uncomment to enable accelerated launch +#export ACCELERATE="True" + ########################################### diff --git a/webui.bat b/webui.bat index a38a28bb..d4d626e2 100644 --- a/webui.bat +++ b/webui.bat @@ -28,15 +28,27 @@ goto :show_stdout_stderr :activate_venv set PYTHON="%~dp0%VENV_DIR%\Scripts\Python.exe" echo venv %PYTHON% +if [%ACCELERATE%] == ["True"] goto :accelerate goto :launch :skip_venv +:accelerate +echo "Checking for accelerate" +set ACCELERATE="%~dp0%VENV_DIR%\Scripts\accelerate.exe" +if EXIST %ACCELERATE% goto :accelerate_launch + :launch %PYTHON% launch.py %* pause exit /b +:accelerate_launch +echo "Accelerating" +%ACCELERATE% launch --num_cpu_threads_per_process=6 launch.py +pause +exit /b + :show_stdout_stderr echo. diff --git a/webui.py b/webui.py index f4f1d74d..2cd3bae9 100644 --- a/webui.py +++ b/webui.py @@ -82,6 +82,7 @@ def initialize(): modules.sd_models.load_model() shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: modules.sd_models.reload_model_weights())) shared.opts.onchange("sd_vae", wrap_queued_call(lambda: modules.sd_vae.reload_vae_weights()), call=False) + shared.opts.onchange("sd_vae_as_default", wrap_queued_call(lambda: modules.sd_vae.reload_vae_weights()), call=False) shared.opts.onchange("sd_hypernetwork", wrap_queued_call(lambda: modules.hypernetworks.hypernetwork.load_hypernetwork(shared.opts.sd_hypernetwork))) shared.opts.onchange("sd_hypernetwork_strength", modules.hypernetworks.hypernetwork.apply_strength) diff --git a/webui.sh b/webui.sh index cc06f344..d89b344b 100755 --- a/webui.sh +++ b/webui.sh @@ -134,7 +134,15 @@ else exit 1 fi -printf "\n%s\n" "${delimiter}" -printf "Launching launch.py..." -printf "\n%s\n" "${delimiter}" -"${python_cmd}" "${LAUNCH_SCRIPT}" "$@" +if [[ ! -z "${ACCELERATE}" ]] && [ ${ACCELERATE}="True" ] && [ -x "$(command -v accelerate)" ] +then + printf "\n%s\n" "${delimiter}" + printf "Accelerating launch.py..." + printf "\n%s\n" "${delimiter}" + accelerate launch --num_cpu_threads_per_process=6 "${LAUNCH_SCRIPT}" "$@" +else + printf "\n%s\n" "${delimiter}" + printf "Launching launch.py..." + printf "\n%s\n" "${delimiter}" + "${python_cmd}" "${LAUNCH_SCRIPT}" "$@" +fi