From 423f22228306ae72d0480e25add9777c3c5d8fdf Mon Sep 17 00:00:00 2001 From: Maiko Tan Date: Sun, 30 Oct 2022 22:46:43 +0800 Subject: feat: add app started callback --- webui.py | 3 +++ 1 file changed, 3 insertions(+) (limited to 'webui.py') diff --git a/webui.py b/webui.py index 29530872..13407e42 100644 --- a/webui.py +++ b/webui.py @@ -23,6 +23,7 @@ import modules.sd_hijack import modules.sd_models import modules.shared as shared import modules.txt2img +import modules.script_callbacks import modules.ui from modules import devices @@ -135,6 +136,8 @@ def webui(): if (launch_api): create_api(app) + modules.script_callbacks.app_started_callback(demo, app) + wait_on_server(demo) sd_samplers.set_samplers() -- cgit v1.2.3 From cb31abcf58ea1f64266e6d821937eed058c35f4d Mon Sep 17 00:00:00 2001 From: Muhammad Rizqi Nur Date: Sun, 30 Oct 2022 21:54:31 +0700 Subject: Settings to select VAE --- modules/sd_models.py | 31 +++++-------- modules/sd_vae.py | 121 +++++++++++++++++++++++++++++++++++++++++++++++++++ modules/shared.py | 8 ++-- webui.py | 5 +++ 4 files changed, 141 insertions(+), 24 deletions(-) create mode 100644 modules/sd_vae.py (limited to 'webui.py') diff --git a/modules/sd_models.py b/modules/sd_models.py index f86dc3ed..91ad4b5e 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -8,7 +8,7 @@ from omegaconf import OmegaConf from ldm.util import instantiate_from_config -from modules import shared, modelloader, devices, script_callbacks +from modules import shared, modelloader, devices, script_callbacks, sd_vae from modules.paths import models_path from modules.sd_hijack_inpainting import do_inpainting_hijack, should_hijack_inpainting @@ -160,12 +160,11 @@ def get_state_dict_from_checkpoint(pl_sd): vae_ignore_keys = {"model_ema.decay", "model_ema.num_updates"} - -def load_model_weights(model, checkpoint_info): +def load_model_weights(model, checkpoint_info, force=False): checkpoint_file = checkpoint_info.filename sd_model_hash = checkpoint_info.hash - if checkpoint_info not in checkpoints_loaded: + if force or checkpoint_info not in checkpoints_loaded: print(f"Loading weights [{sd_model_hash}] from {checkpoint_file}") pl_sd = torch.load(checkpoint_file, map_location=shared.weight_load_location) @@ -186,17 +185,7 @@ def load_model_weights(model, checkpoint_info): devices.dtype = torch.float32 if shared.cmd_opts.no_half else torch.float16 devices.dtype_vae = torch.float32 if shared.cmd_opts.no_half or shared.cmd_opts.no_half_vae else torch.float16 - vae_file = os.path.splitext(checkpoint_file)[0] + ".vae.pt" - - if not os.path.exists(vae_file) and shared.cmd_opts.vae_path is not None: - vae_file = shared.cmd_opts.vae_path - - if os.path.exists(vae_file): - print(f"Loading VAE weights from: {vae_file}") - vae_ckpt = torch.load(vae_file, map_location=shared.weight_load_location) - vae_dict = {k: v for k, v in vae_ckpt["state_dict"].items() if k[0:4] != "loss" and k not in vae_ignore_keys} - model.first_stage_model.load_state_dict(vae_dict) - + sd_vae.load_vae(model, checkpoint_file) model.first_stage_model.to(devices.dtype_vae) if shared.opts.sd_checkpoint_cache > 0: @@ -213,7 +202,7 @@ def load_model_weights(model, checkpoint_info): model.sd_checkpoint_info = checkpoint_info -def load_model(checkpoint_info=None): +def load_model(checkpoint_info=None, force=False): from modules import lowvram, sd_hijack checkpoint_info = checkpoint_info or select_checkpoint() @@ -234,7 +223,7 @@ def load_model(checkpoint_info=None): do_inpainting_hijack() sd_model = instantiate_from_config(sd_config.model) - load_model_weights(sd_model, checkpoint_info) + load_model_weights(sd_model, checkpoint_info, force=force) if shared.cmd_opts.lowvram or shared.cmd_opts.medvram: lowvram.setup_for_low_vram(sd_model, shared.cmd_opts.medvram) @@ -252,16 +241,16 @@ def load_model(checkpoint_info=None): return sd_model -def reload_model_weights(sd_model, info=None): +def reload_model_weights(sd_model, info=None, force=False): from modules import lowvram, devices, sd_hijack checkpoint_info = info or select_checkpoint() - if sd_model.sd_model_checkpoint == checkpoint_info.filename: + if sd_model.sd_model_checkpoint == checkpoint_info.filename and not force: return if sd_model.sd_checkpoint_info.config != checkpoint_info.config or should_hijack_inpainting(checkpoint_info) != should_hijack_inpainting(sd_model.sd_checkpoint_info): checkpoints_loaded.clear() - load_model(checkpoint_info) + load_model(checkpoint_info, force=force) return shared.sd_model if shared.cmd_opts.lowvram or shared.cmd_opts.medvram: @@ -271,7 +260,7 @@ def reload_model_weights(sd_model, info=None): sd_hijack.model_hijack.undo_hijack(sd_model) - load_model_weights(sd_model, checkpoint_info) + load_model_weights(sd_model, checkpoint_info, force=force) sd_hijack.model_hijack.hijack(sd_model) script_callbacks.model_loaded_callback(sd_model) diff --git a/modules/sd_vae.py b/modules/sd_vae.py new file mode 100644 index 00000000..82764e55 --- /dev/null +++ b/modules/sd_vae.py @@ -0,0 +1,121 @@ +import torch +import os +from collections import namedtuple +from modules import shared, devices +from modules.paths import models_path +import glob + +model_dir = "Stable-diffusion" +model_path = os.path.abspath(os.path.join(models_path, model_dir)) +vae_dir = "VAE" +vae_path = os.path.abspath(os.path.join(models_path, vae_dir)) + +vae_ignore_keys = {"model_ema.decay", "model_ema.num_updates"} +default_vae_dict = {"auto": "auto", "None": "None"} +default_vae_list = ["auto", "None"] +default_vae_values = [default_vae_dict[x] for x in default_vae_list] +vae_dict = dict(default_vae_dict) +vae_list = list(default_vae_list) +first_load = True + +def get_filename(filepath): + return os.path.splitext(os.path.basename(filepath))[0] + +def refresh_vae_list(vae_path=vae_path, model_path=model_path): + global vae_dict, vae_list + res = {} + candidates = [ + *glob.iglob(os.path.join(model_path, '**/*.vae.pt'), recursive=True), + *glob.iglob(os.path.join(model_path, '**/*.vae.ckpt'), recursive=True), + *glob.iglob(os.path.join(vae_path, '**/*.pt'), recursive=True), + *glob.iglob(os.path.join(vae_path, '**/*.ckpt'), recursive=True) + ] + if shared.cmd_opts.vae_path is not None and os.path.isfile(shared.cmd_opts.vae_path): + candidates.append(shared.cmd_opts.vae_path) + for filepath in candidates: + name = get_filename(filepath) + res[name] = filepath + vae_list.clear() + vae_list.extend(default_vae_list) + vae_list.extend(list(res.keys())) + vae_dict.clear() + vae_dict.update(default_vae_dict) + vae_dict.update(res) + return vae_list + +def load_vae(model, checkpoint_file, vae_file="auto"): + global first_load, vae_dict, vae_list + # save_settings = False + + # if vae_file argument is provided, it takes priority + if vae_file and vae_file not in default_vae_list: + if not os.path.isfile(vae_file): + vae_file = "auto" + # save_settings = True + print("VAE provided as function argument doesn't exist") + # for the first load, if vae-path is provided, it takes priority and failure is reported + if first_load and shared.cmd_opts.vae_path is not None: + if os.path.isfile(shared.cmd_opts.vae_path): + vae_file = shared.cmd_opts.vae_path + # save_settings = True + # print("Using VAE provided as command line argument") + 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") + # 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): + vae_file = shared.cmd_opts.vae_path + print("Using VAE provided as command line argument") + # if still not found, try look for ".vae.pt" beside model + model_path = os.path.splitext(checkpoint_file)[0] + if vae_file == "auto": + vae_file_try = model_path + ".vae.pt" + if os.path.isfile(vae_file_try): + vae_file = vae_file_try + print("Using VAE found beside selected model") + # if still not found, try look for ".vae.ckpt" beside model + if vae_file == "auto": + vae_file_try = model_path + ".vae.ckpt" + if os.path.isfile(vae_file_try): + vae_file = vae_file_try + print("Using VAE found beside selected model") + # No more fallbacks for auto + if vae_file == "auto": + vae_file = None + # Last check, just because + if vae_file and not os.path.exists(vae_file): + vae_file = None + + if vae_file: + print(f"Loading VAE weights from: {vae_file}") + vae_ckpt = torch.load(vae_file, map_location=shared.weight_load_location) + vae_dict_1 = {k: v for k, v in vae_ckpt["state_dict"].items() if k[0:4] != "loss" and k not in vae_ignore_keys} + model.first_stage_model.load_state_dict(vae_dict_1) + + # If vae used is not in dict, update it + # It will be removed on refresh though + if vae_file is not None: + vae_opt = get_filename(vae_file) + if vae_opt not in vae_dict: + vae_dict[vae_opt] = vae_file + vae_list.append(vae_opt) + + """ + # Save current VAE to VAE settings, maybe? will it work? + if save_settings: + if vae_file is None: + vae_opt = "None" + + # shared.opts.sd_vae = vae_opt + """ + + first_load = False + model.first_stage_model.to(devices.dtype_vae) diff --git a/modules/shared.py b/modules/shared.py index e4f163c1..06440ac4 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -14,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 +from modules import sd_samplers, sd_models, localization, sd_vae from modules.hypernetworks import hypernetwork from modules.paths import models_path, script_path, sd_path @@ -295,6 +295,7 @@ options_templates.update(options_section(('training', "Training"), { 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": list(sd_vae.vae_list)}, refresh=sd_vae.refresh_vae_list), "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}), @@ -407,11 +408,12 @@ class Options: if bad_settings > 0: print(f"The program is likely to not work with bad settings.\nSettings file: {filename}\nEither fix the file, or delete it and restart.", file=sys.stderr) - def onchange(self, key, func): + def onchange(self, key, func, call=True): item = self.data_labels.get(key) item.onchange = func - func() + if call: + func() def dumpjson(self): d = {k: self.data.get(k, self.data_labels.get(k).default) for k in self.data_labels.keys()} diff --git a/webui.py b/webui.py index 29530872..27949f3d 100644 --- a/webui.py +++ b/webui.py @@ -21,6 +21,7 @@ import modules.paths import modules.scripts import modules.sd_hijack import modules.sd_models +import modules.sd_vae import modules.shared as shared import modules.txt2img @@ -74,8 +75,12 @@ def initialize(): modules.scripts.load_scripts() + modules.sd_vae.refresh_vae_list() modules.sd_models.load_model() shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(shared.sd_model))) + # I don't know what needs to be done to only reload VAE, with all those hijacks callbacks, and lowvram, + # so for now this reloads the whole model too, and no cache + shared.opts.onchange("sd_vae", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(shared.sd_model, force=True)), 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) -- cgit v1.2.3 From 910a097ae2ed78a62101951f1b87137f9e1baaea Mon Sep 17 00:00:00 2001 From: AUTOMATIC <16777216c@gmail.com> Date: Mon, 31 Oct 2022 17:36:45 +0300 Subject: add initial version of the extensions tab fix broken Restart Gradio button --- javascript/extensions.js | 24 +++++ modules/extensions.py | 83 +++++++++++++++ modules/generation_parameters_copypaste.py | 5 + modules/scripts.py | 21 +--- modules/shared.py | 10 +- modules/ui.py | 16 ++- modules/ui_extensions.py | 162 +++++++++++++++++++++++++++++ style.css | 22 +++- webui.py | 20 ++-- 9 files changed, 333 insertions(+), 30 deletions(-) create mode 100644 javascript/extensions.js create mode 100644 modules/extensions.py create mode 100644 modules/ui_extensions.py (limited to 'webui.py') diff --git a/javascript/extensions.js b/javascript/extensions.js new file mode 100644 index 00000000..86f5336d --- /dev/null +++ b/javascript/extensions.js @@ -0,0 +1,24 @@ + +function extensions_apply(_, _){ + disable = [] + update = [] + gradioApp().querySelectorAll('#extensions input[type="checkbox"]').forEach(function(x){ + if(x.name.startsWith("enable_") && ! x.checked) + disable.push(x.name.substr(7)) + + if(x.name.startsWith("update_") && x.checked) + update.push(x.name.substr(7)) + }) + + restart_reload() + + return [JSON.stringify(disable), JSON.stringify(update)] +} + +function extensions_check(){ + gradioApp().querySelectorAll('#extensions .extension_status').forEach(function(x){ + x.innerHTML = "Loading..." + }) + + return [] +} \ No newline at end of file diff --git a/modules/extensions.py b/modules/extensions.py new file mode 100644 index 00000000..8d6ae848 --- /dev/null +++ b/modules/extensions.py @@ -0,0 +1,83 @@ +import os +import sys +import traceback + +import git + +from modules import paths, shared + + +extensions = [] +extensions_dir = os.path.join(paths.script_path, "extensions") + + +def active(): + return [x for x in extensions if x.enabled] + + +class Extension: + def __init__(self, name, path, enabled=True): + self.name = name + self.path = path + self.enabled = enabled + self.status = '' + self.can_update = False + + repo = None + try: + if os.path.exists(os.path.join(path, ".git")): + repo = git.Repo(path) + except Exception: + print(f"Error reading github repository info from {path}:", file=sys.stderr) + print(traceback.format_exc(), file=sys.stderr) + + if repo is None or repo.bare: + self.remote = None + else: + self.remote = next(repo.remote().urls, None) + self.status = 'unknown' + + def list_files(self, subdir, extension): + from modules import scripts + + dirpath = os.path.join(self.path, subdir) + if not os.path.isdir(dirpath): + return [] + + res = [] + for filename in sorted(os.listdir(dirpath)): + res.append(scripts.ScriptFile(dirpath, filename, os.path.join(dirpath, filename))) + + res = [x for x in res if os.path.splitext(x.path)[1].lower() == extension and os.path.isfile(x.path)] + + return res + + def check_updates(self): + repo = git.Repo(self.path) + for fetch in repo.remote().fetch("--dry-run"): + if fetch.flags != fetch.HEAD_UPTODATE: + self.can_update = True + self.status = "behind" + return + + self.can_update = False + self.status = "latest" + + def pull(self): + repo = git.Repo(self.path) + repo.remotes.origin.pull() + + +def list_extensions(): + extensions.clear() + + 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 + + extension = Extension(name=dirname, path=path, enabled=dirname not in shared.opts.disabled_extensions) + extensions.append(extension) diff --git a/modules/generation_parameters_copypaste.py b/modules/generation_parameters_copypaste.py index df70c728..985ec95e 100644 --- a/modules/generation_parameters_copypaste.py +++ b/modules/generation_parameters_copypaste.py @@ -17,6 +17,11 @@ paste_fields = {} bind_list = [] +def reset(): + paste_fields.clear() + bind_list.clear() + + def quote(text): if ',' not in str(text): return text diff --git a/modules/scripts.py b/modules/scripts.py index 96e44bfd..533db45c 100644 --- a/modules/scripts.py +++ b/modules/scripts.py @@ -7,7 +7,7 @@ import modules.ui as ui import gradio as gr from modules.processing import StableDiffusionProcessing -from modules import shared, paths, script_callbacks +from modules import shared, paths, script_callbacks, extensions AlwaysVisible = object() @@ -107,17 +107,8 @@ def list_scripts(scriptdirname, extension): for filename in sorted(os.listdir(basedir)): scripts_list.append(ScriptFile(paths.script_path, filename, os.path.join(basedir, filename))) - extdir = os.path.join(paths.script_path, "extensions") - if os.path.exists(extdir): - for dirname in sorted(os.listdir(extdir)): - dirpath = os.path.join(extdir, dirname) - scriptdirpath = os.path.join(dirpath, scriptdirname) - - if not os.path.isdir(scriptdirpath): - continue - - for filename in sorted(os.listdir(scriptdirpath)): - scripts_list.append(ScriptFile(dirpath, filename, os.path.join(scriptdirpath, filename))) + for ext in extensions.active(): + scripts_list += ext.list_files(scriptdirname, extension) scripts_list = [x for x in scripts_list if os.path.splitext(x.path)[1].lower() == extension and os.path.isfile(x.path)] @@ -127,11 +118,7 @@ def list_scripts(scriptdirname, extension): def list_files_with_name(filename): res = [] - dirs = [paths.script_path] - - extdir = os.path.join(paths.script_path, "extensions") - if os.path.exists(extdir): - dirs += [os.path.join(extdir, d) for d in sorted(os.listdir(extdir))] + dirs = [paths.script_path] + [ext.path for ext in extensions.active()] for dirpath in dirs: if not os.path.isdir(dirpath): diff --git a/modules/shared.py b/modules/shared.py index e4f163c1..cce87081 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -132,6 +132,7 @@ class State: current_image = None current_image_sampling_step = 0 textinfo = None + need_restart = False def skip(self): self.skipped = True @@ -354,6 +355,12 @@ options_templates.update(options_section(('sampler-params', "Sampler parameters" 'eta_noise_seed_delta': OptionInfo(0, "Eta noise seed delta", gr.Number, {"precision": 0}), })) +options_templates.update(options_section((None, "Hidden options"), { + "disabled_extensions": OptionInfo([], "Disable those extensions"), +})) + +options_templates.update() + class Options: data = None @@ -365,8 +372,9 @@ class Options: def __setattr__(self, key, value): if self.data is not None: - if key in self.data: + if key in self.data or key in self.data_labels: self.data[key] = value + return return super(Options, self).__setattr__(key, value) diff --git a/modules/ui.py b/modules/ui.py index 5055ca64..2c15abb7 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -19,7 +19,7 @@ import numpy as np from PIL import Image, PngImagePlugin -from modules import sd_hijack, sd_models, localization, script_callbacks +from modules import sd_hijack, sd_models, localization, script_callbacks, ui_extensions from modules.paths import script_path from modules.shared import opts, cmd_opts, restricted_opts @@ -671,6 +671,7 @@ def create_ui(wrap_gradio_gpu_call): import modules.img2img import modules.txt2img + parameters_copypaste.reset() with gr.Blocks(analytics_enabled=False) as txt2img_interface: txt2img_prompt, roll, txt2img_prompt_style, txt2img_negative_prompt, txt2img_prompt_style2, submit, _, _, txt2img_prompt_style_apply, txt2img_save_style, txt2img_paste, token_counter, token_button = create_toprow(is_img2img=False) @@ -1511,8 +1512,9 @@ def create_ui(wrap_gradio_gpu_call): column = None with gr.Row(elem_id="settings").style(equal_height=False): for i, (k, item) in enumerate(opts.data_labels.items()): + section_must_be_skipped = item.section[0] is None - if previous_section != item.section: + if previous_section != item.section and not section_must_be_skipped: if cols_displayed < settings_cols and (items_displayed >= items_per_col or previous_section is None): if column is not None: column.__exit__() @@ -1531,6 +1533,8 @@ def create_ui(wrap_gradio_gpu_call): if k in quicksettings_names and not shared.cmd_opts.freeze_settings: quicksettings_list.append((i, k, item)) components.append(dummy_component) + elif section_must_be_skipped: + components.append(dummy_component) else: component = create_setting_component(k) component_dict[k] = component @@ -1572,9 +1576,10 @@ def create_ui(wrap_gradio_gpu_call): def request_restart(): shared.state.interrupt() - settings_interface.gradio_ref.do_restart = True + shared.state.need_restart = True restart_gradio.click( + fn=request_restart, inputs=[], outputs=[], @@ -1612,14 +1617,15 @@ def create_ui(wrap_gradio_gpu_call): interfaces += script_callbacks.ui_tabs_callback() interfaces += [(settings_interface, "Settings", "settings")] + extensions_interface = ui_extensions.create_ui() + interfaces += [(extensions_interface, "Extensions", "extensions")] + with gr.Blocks(css=css, analytics_enabled=False, title="Stable Diffusion") as demo: with gr.Row(elem_id="quicksettings"): for i, k, item in quicksettings_list: component = create_setting_component(k, is_quicksettings=True) component_dict[k] = component - settings_interface.gradio_ref = demo - parameters_copypaste.integrate_settings_paste_fields(component_dict) parameters_copypaste.run_bind() diff --git a/modules/ui_extensions.py b/modules/ui_extensions.py new file mode 100644 index 00000000..b7d747dc --- /dev/null +++ b/modules/ui_extensions.py @@ -0,0 +1,162 @@ +import json +import os.path +import shutil +import sys +import time +import traceback + +import git + +import gradio as gr +import html + +from modules import extensions, shared, paths + + +def apply_and_restart(disable_list, update_list): + disabled = json.loads(disable_list) + assert type(disabled) == list, f"wrong disable_list data for apply_and_restart: {disable_list}" + + update = json.loads(update_list) + assert type(update) == list, f"wrong update_list data for apply_and_restart: {update_list}" + + update = set(update) + + for ext in extensions.extensions: + if ext.name not in update: + continue + + try: + ext.pull() + except Exception: + print(f"Error pulling updates for {ext.name}:", file=sys.stderr) + print(traceback.format_exc(), file=sys.stderr) + + shared.opts.disabled_extensions = disabled + shared.opts.save(shared.config_filename) + + shared.state.interrupt() + shared.state.need_restart = True + + +def check_updates(): + for ext in extensions.extensions: + if ext.remote is None: + continue + + try: + ext.check_updates() + except Exception: + print(f"Error checking updates for {ext.name}:", file=sys.stderr) + print(traceback.format_exc(), file=sys.stderr) + + return extension_table() + + +def extension_table(): + code = f""" + + + + + + + + + + """ + + for ext in extensions.extensions: + if ext.can_update: + ext_status = f"""""" + else: + ext_status = ext.status + + code += f""" + + + + {ext_status} + + """ + + code += """ + +
ExtensionURLUpdate
{html.escape(ext.remote or '')}
+ """ + + return code + + +def install_extension_from_url(dirname, url): + assert url, 'No URL specified' + + if dirname is None or dirname == "": + *parts, last_part = url.split('/') + last_part = last_part.replace(".git", "") + + dirname = last_part + + target_dir = os.path.join(extensions.extensions_dir, dirname) + assert not os.path.exists(target_dir), f'Extension directory already exists: {target_dir}' + + assert len([x for x in extensions.extensions if x.remote == url]) == 0, 'Extension with this URL is already installed' + + tmpdir = os.path.join(paths.script_path, "tmp", dirname) + + try: + shutil.rmtree(tmpdir, True) + + repo = git.Repo.clone_from(url, tmpdir) + repo.remote().fetch() + + os.rename(tmpdir, target_dir) + + extensions.list_extensions() + return [extension_table(), html.escape(f"Installed into {target_dir}. Use Installed tab to restart.")] + finally: + shutil.rmtree(tmpdir, True) + + +def create_ui(): + import modules.ui + + with gr.Blocks(analytics_enabled=False) as ui: + with gr.Tabs(elem_id="tabs_extensions") as tabs: + with gr.TabItem("Installed"): + extensions_disabled_list = gr.Text(elem_id="extensions_disabled_list", visible=False) + extensions_update_list = gr.Text(elem_id="extensions_update_list", visible=False) + + with gr.Row(): + apply = gr.Button(value="Apply and restart UI", variant="primary") + check = gr.Button(value="Check for updates") + + extensions_table = gr.HTML(lambda: extension_table()) + + apply.click( + fn=apply_and_restart, + _js="extensions_apply", + inputs=[extensions_disabled_list, extensions_update_list], + outputs=[], + ) + + check.click( + fn=check_updates, + _js="extensions_check", + inputs=[], + outputs=[extensions_table], + ) + + with gr.TabItem("Install from URL"): + install_url = gr.Text(label="URL for extension's git repository") + install_dirname = gr.Text(label="Local directory name", placeholder="Leave empty for auto") + intall_button = gr.Button(value="Install", variant="primary") + intall_result = gr.HTML(elem_id="extension_install_result") + + intall_button.click( + fn=modules.ui.wrap_gradio_call(install_extension_from_url, extra_outputs=[gr.update()]), + inputs=[install_dirname, install_url], + outputs=[extensions_table, intall_result], + ) + + return ui diff --git a/style.css b/style.css index 8b2211b1..859c3933 100644 --- a/style.css +++ b/style.css @@ -530,6 +530,26 @@ img2maskimg, #img2maskimg > .h-60, #img2maskimg > .h-60 > div, #img2maskimg > .h min-height: 480px !important; } +/* Extensions */ + +#extensions{ + border-collapse: collapse; +} + +#extensions td, #extensions th{ + border: 1px solid #ccc; + padding: 0.25em 0.5em; +} + +#extensions input[type="checkbox"]{ + margin-right: 0.5em; +} + +#tab_extensions button{ + max-width: 16em; +} + + /* The following handles localization for right-to-left (RTL) languages like Arabic. The rtl media type will only be activated by the logic in javascript/localization.js. If you change anything above, you need to make sure it is RTL compliant by just running @@ -607,4 +627,4 @@ Then, you will need to add the RTL counterpart only if needed in the rtl section right: unset; left: 0.5em; } -} +} \ No newline at end of file diff --git a/webui.py b/webui.py index 29530872..ad2eb236 100644 --- a/webui.py +++ b/webui.py @@ -9,7 +9,7 @@ from fastapi.middleware.gzip import GZipMiddleware from modules.paths import script_path -from modules import devices, sd_samplers, upscaler +from modules import devices, sd_samplers, upscaler, extensions import modules.codeformer_model as codeformer import modules.extras import modules.face_restoration @@ -60,6 +60,11 @@ def wrap_gradio_gpu_call(func, extra_outputs=None): def initialize(): + extensions.list_extensions() + #for ext in extensions.extensions: + # print(ext.name, ext.path, ext.enabled, ext.remote) + #exit() + if cmd_opts.ui_debug_mode: shared.sd_upscalers = upscaler.UpscalerLanczos().scalers modules.scripts.load_scripts() @@ -92,15 +97,18 @@ def create_api(app): api = Api(app, queue_lock) return api + def wait_on_server(demo=None): while 1: time.sleep(0.5) - if demo and getattr(demo, 'do_restart', False): + if shared.state.need_restart: + shared.state.need_restart = False time.sleep(0.5) demo.close() time.sleep(0.5) break + def api_only(): initialize() @@ -132,14 +140,16 @@ def webui(): app.add_middleware(GZipMiddleware, minimum_size=1000) - if (launch_api): + if launch_api: create_api(app) wait_on_server(demo) sd_samplers.set_samplers() - print('Reloading Custom Scripts') + print('Reloading extensions') + extensions.list_extensions() + print('Reloading custom scripts') modules.scripts.reload_scripts() print('Reloading modules: modules.ui') importlib.reload(modules.ui) @@ -148,8 +158,6 @@ def webui(): print('Restarting Gradio') - -task = [] if __name__ == "__main__": if cmd_opts.nowebui: api_only() -- cgit v1.2.3 From 58cc03edd0fe8c7e64297bcfe51111caaafabfd7 Mon Sep 17 00:00:00 2001 From: AUTOMATIC <16777216c@gmail.com> Date: Mon, 31 Oct 2022 18:40:47 +0300 Subject: fix scripts I broke with the extension tab changes --- modules/extensions.py | 2 +- webui.py | 3 --- 2 files changed, 1 insertion(+), 4 deletions(-) (limited to 'webui.py') diff --git a/modules/extensions.py b/modules/extensions.py index 8d6ae848..897af96e 100644 --- a/modules/extensions.py +++ b/modules/extensions.py @@ -46,7 +46,7 @@ class Extension: res = [] for filename in sorted(os.listdir(dirpath)): - res.append(scripts.ScriptFile(dirpath, filename, os.path.join(dirpath, filename))) + res.append(scripts.ScriptFile(self.path, filename, os.path.join(dirpath, filename))) res = [x for x in res if os.path.splitext(x.path)[1].lower() == extension and os.path.isfile(x.path)] diff --git a/webui.py b/webui.py index ad2eb236..6ff95dc4 100644 --- a/webui.py +++ b/webui.py @@ -61,9 +61,6 @@ def wrap_gradio_gpu_call(func, extra_outputs=None): def initialize(): extensions.list_extensions() - #for ext in extensions.extensions: - # print(ext.name, ext.path, ext.enabled, ext.remote) - #exit() if cmd_opts.ui_debug_mode: shared.sd_upscalers = upscaler.UpscalerLanczos().scalers -- cgit v1.2.3 From af758e97fa2c4c853042f121af4e974be01e6696 Mon Sep 17 00:00:00 2001 From: Jairo Correa Date: Tue, 1 Nov 2022 04:01:49 -0300 Subject: Unload sd_model before loading the other --- modules/lowvram.py | 21 +++++++++++++-------- modules/processing.py | 3 +++ modules/sd_hijack.py | 4 ++++ modules/sd_models.py | 14 +++++++++++++- webui.py | 2 +- 5 files changed, 34 insertions(+), 10 deletions(-) (limited to 'webui.py') diff --git a/modules/lowvram.py b/modules/lowvram.py index f327c3df..a4652cb1 100644 --- a/modules/lowvram.py +++ b/modules/lowvram.py @@ -38,13 +38,18 @@ def setup_for_low_vram(sd_model, use_medvram): # see below for register_forward_pre_hook; # first_stage_model does not use forward(), it uses encode/decode, so register_forward_pre_hook is # useless here, and we just replace those methods - def first_stage_model_encode_wrap(self, encoder, x): - send_me_to_gpu(self, None) - return encoder(x) - def first_stage_model_decode_wrap(self, decoder, z): - send_me_to_gpu(self, None) - return decoder(z) + first_stage_model = sd_model.first_stage_model + first_stage_model_encode = sd_model.first_stage_model.encode + first_stage_model_decode = sd_model.first_stage_model.decode + + def first_stage_model_encode_wrap(x): + send_me_to_gpu(first_stage_model, None) + return first_stage_model_encode(x) + + def first_stage_model_decode_wrap(z): + send_me_to_gpu(first_stage_model, None) + return first_stage_model_decode(z) # remove three big modules, cond, first_stage, and unet from the model and then # send the model to GPU. Then put modules back. the modules will be in CPU. @@ -56,8 +61,8 @@ def setup_for_low_vram(sd_model, use_medvram): # register hooks for those the first two models sd_model.cond_stage_model.transformer.register_forward_pre_hook(send_me_to_gpu) sd_model.first_stage_model.register_forward_pre_hook(send_me_to_gpu) - sd_model.first_stage_model.encode = lambda x, en=sd_model.first_stage_model.encode: first_stage_model_encode_wrap(sd_model.first_stage_model, en, x) - sd_model.first_stage_model.decode = lambda z, de=sd_model.first_stage_model.decode: first_stage_model_decode_wrap(sd_model.first_stage_model, de, z) + sd_model.first_stage_model.encode = first_stage_model_encode_wrap + sd_model.first_stage_model.decode = first_stage_model_decode_wrap parents[sd_model.cond_stage_model.transformer] = sd_model.cond_stage_model if use_medvram: diff --git a/modules/processing.py b/modules/processing.py index b1df4918..57d3a523 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -597,6 +597,9 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: if p.scripts is not None: p.scripts.postprocess(p, res) + p.sd_model = None + p.sampler = None + return res diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py index 0f10828e..bc49d235 100644 --- a/modules/sd_hijack.py +++ b/modules/sd_hijack.py @@ -94,6 +94,10 @@ class StableDiffusionModelHijack: if type(model_embeddings.token_embedding) == EmbeddingsWithFixes: model_embeddings.token_embedding = model_embeddings.token_embedding.wrapped + self.layers = None + self.circular_enabled = False + self.clip = None + def apply_circular(self, enable): if self.circular_enabled == enable: return diff --git a/modules/sd_models.py b/modules/sd_models.py index f86dc3ed..90007da3 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -1,6 +1,7 @@ import collections import os.path import sys +import gc from collections import namedtuple import torch import re @@ -220,6 +221,12 @@ def load_model(checkpoint_info=None): if checkpoint_info.config != shared.cmd_opts.config: print(f"Loading config from: {checkpoint_info.config}") + if shared.sd_model: + sd_hijack.model_hijack.undo_hijack(shared.sd_model) + shared.sd_model = None + gc.collect() + devices.torch_gc() + sd_config = OmegaConf.load(checkpoint_info.config) if should_hijack_inpainting(checkpoint_info): @@ -233,6 +240,7 @@ def load_model(checkpoint_info=None): checkpoint_info = checkpoint_info._replace(config=checkpoint_info.config.replace(".yaml", "-inpainting.yaml")) do_inpainting_hijack() + sd_model = instantiate_from_config(sd_config.model) load_model_weights(sd_model, checkpoint_info) @@ -252,14 +260,18 @@ def load_model(checkpoint_info=None): return sd_model -def reload_model_weights(sd_model, info=None): +def reload_model_weights(sd_model=None, info=None): from modules import lowvram, devices, sd_hijack checkpoint_info = info or select_checkpoint() + if not sd_model: + sd_model = shared.sd_model + if sd_model.sd_model_checkpoint == checkpoint_info.filename: return if sd_model.sd_checkpoint_info.config != checkpoint_info.config or should_hijack_inpainting(checkpoint_info) != should_hijack_inpainting(sd_model.sd_checkpoint_info): + del sd_model checkpoints_loaded.clear() load_model(checkpoint_info) return shared.sd_model diff --git a/webui.py b/webui.py index 6ff95dc4..9c393e55 100644 --- a/webui.py +++ b/webui.py @@ -77,7 +77,7 @@ def initialize(): modules.scripts.load_scripts() modules.sd_models.load_model() - shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(shared.sd_model))) + shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: modules.sd_models.reload_model_weights())) 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) -- cgit v1.2.3 From 056f06d3738c267b1014e6e8e1ef5bd97af1fb45 Mon Sep 17 00:00:00 2001 From: Muhammad Rizqi Nur Date: Wed, 2 Nov 2022 12:51:46 +0700 Subject: Reload VAE without reloading sd checkpoint --- modules/sd_models.py | 15 ++++---- modules/sd_vae.py | 97 ++++++++++++++++++++++++++++++++++++++++++++++++---- webui.py | 4 +-- 3 files changed, 98 insertions(+), 18 deletions(-) (limited to 'webui.py') diff --git a/modules/sd_models.py b/modules/sd_models.py index 6ab85b65..883639d1 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -159,15 +159,13 @@ def get_state_dict_from_checkpoint(pl_sd): return pl_sd -vae_ignore_keys = {"model_ema.decay", "model_ema.num_updates"} - def load_model_weights(model, checkpoint_info, vae_file="auto"): checkpoint_file = checkpoint_info.filename sd_model_hash = checkpoint_info.hash vae_file = sd_vae.resolve_vae(checkpoint_file, vae_file=vae_file) - checkpoint_key = (checkpoint_info, vae_file) + checkpoint_key = checkpoint_info if checkpoint_key not in checkpoints_loaded: print(f"Loading weights [{sd_model_hash}] from {checkpoint_file}") @@ -190,13 +188,12 @@ def load_model_weights(model, checkpoint_info, vae_file="auto"): devices.dtype = torch.float32 if shared.cmd_opts.no_half else torch.float16 devices.dtype_vae = torch.float32 if shared.cmd_opts.no_half or shared.cmd_opts.no_half_vae else torch.float16 - sd_vae.load_vae(model, vae_file) - model.first_stage_model.to(devices.dtype_vae) - if shared.opts.sd_checkpoint_cache > 0: + # if PR #4035 were to get merged, restore base VAE first before caching checkpoints_loaded[checkpoint_key] = model.state_dict().copy() while len(checkpoints_loaded) > shared.opts.sd_checkpoint_cache: checkpoints_loaded.popitem(last=False) # LRU + else: vae_name = sd_vae.get_filename(vae_file) print(f"Loading weights [{sd_model_hash}] with {vae_name} VAE from cache") @@ -207,6 +204,8 @@ def load_model_weights(model, checkpoint_info, vae_file="auto"): model.sd_model_checkpoint = checkpoint_file model.sd_checkpoint_info = checkpoint_info + sd_vae.load_vae(model, vae_file) + def load_model(checkpoint_info=None): from modules import lowvram, sd_hijack @@ -254,14 +253,14 @@ def load_model(checkpoint_info=None): return sd_model -def reload_model_weights(sd_model=None, info=None, force=False): +def reload_model_weights(sd_model=None, info=None): from modules import lowvram, devices, sd_hijack checkpoint_info = info or select_checkpoint() if not sd_model: sd_model = shared.sd_model - if sd_model.sd_model_checkpoint == checkpoint_info.filename and not force: + if sd_model.sd_model_checkpoint == checkpoint_info.filename: return if sd_model.sd_checkpoint_info.config != checkpoint_info.config or should_hijack_inpainting(checkpoint_info) != should_hijack_inpainting(sd_model.sd_checkpoint_info): diff --git a/modules/sd_vae.py b/modules/sd_vae.py index e9239326..78e14e8a 100644 --- a/modules/sd_vae.py +++ b/modules/sd_vae.py @@ -1,26 +1,65 @@ import torch import os from collections import namedtuple -from modules import shared, devices +from modules import shared, devices, script_callbacks from modules.paths import models_path import glob + model_dir = "Stable-diffusion" model_path = os.path.abspath(os.path.join(models_path, model_dir)) vae_dir = "VAE" vae_path = os.path.abspath(os.path.join(models_path, vae_dir)) + vae_ignore_keys = {"model_ema.decay", "model_ema.num_updates"} + + default_vae_dict = {"auto": "auto", "None": "None"} default_vae_list = ["auto", "None"] + + default_vae_values = [default_vae_dict[x] for x in default_vae_list] vae_dict = dict(default_vae_dict) vae_list = list(default_vae_list) first_load = True + +base_vae = None +loaded_vae_file = None +checkpoint_info = None + + +def get_base_vae(model): + if base_vae is not None and checkpoint_info == model.sd_checkpoint_info and model: + return base_vae + return None + + +def store_base_vae(model): + global base_vae, checkpoint_info + if checkpoint_info != model.sd_checkpoint_info: + base_vae = model.first_stage_model.state_dict().copy() + checkpoint_info = model.sd_checkpoint_info + + +def delete_base_vae(): + global base_vae, checkpoint_info + base_vae = None + checkpoint_info = None + + +def restore_base_vae(model): + global base_vae, checkpoint_info + if base_vae is not None and checkpoint_info == model.sd_checkpoint_info: + load_vae_dict(model, base_vae) + delete_base_vae() + + def get_filename(filepath): return os.path.splitext(os.path.basename(filepath))[0] + def refresh_vae_list(vae_path=vae_path, model_path=model_path): global vae_dict, vae_list res = {} @@ -43,6 +82,7 @@ def refresh_vae_list(vae_path=vae_path, model_path=model_path): vae_dict.update(res) return vae_list + def resolve_vae(checkpoint_file, vae_file="auto"): global first_load, vae_dict, vae_list # save_settings = False @@ -96,24 +136,26 @@ def resolve_vae(checkpoint_file, vae_file="auto"): return vae_file -def load_vae(model, vae_file): - global first_load, vae_dict, vae_list + +def load_vae(model, vae_file=None): + global first_load, vae_dict, vae_list, loaded_vae_file # save_settings = False if vae_file: print(f"Loading VAE weights from: {vae_file}") vae_ckpt = torch.load(vae_file, map_location=shared.weight_load_location) vae_dict_1 = {k: v for k, v in vae_ckpt["state_dict"].items() if k[0:4] != "loss" and k not in vae_ignore_keys} - model.first_stage_model.load_state_dict(vae_dict_1) + load_vae_dict(model, vae_dict_1) - # If vae used is not in dict, update it - # It will be removed on refresh though - if vae_file is not None: + # If vae used is not in dict, update it + # It will be removed on refresh though vae_opt = get_filename(vae_file) if vae_opt not in vae_dict: vae_dict[vae_opt] = vae_file vae_list.append(vae_opt) + loaded_vae_file = vae_file + """ # Save current VAE to VAE settings, maybe? will it work? if save_settings: @@ -124,4 +166,45 @@ def load_vae(model, vae_file): """ first_load = False + + +# don't call this from outside +def load_vae_dict(model, vae_dict_1=None): + if vae_dict_1: + store_base_vae(model) + model.first_stage_model.load_state_dict(vae_dict_1) + else: + restore_base_vae() model.first_stage_model.to(devices.dtype_vae) + + +def reload_vae_weights(sd_model=None, vae_file="auto"): + from modules import lowvram, devices, sd_hijack + + if not sd_model: + sd_model = shared.sd_model + + checkpoint_info = sd_model.sd_checkpoint_info + checkpoint_file = checkpoint_info.filename + vae_file = resolve_vae(checkpoint_file, vae_file=vae_file) + + if loaded_vae_file == vae_file: + return + + if shared.cmd_opts.lowvram or shared.cmd_opts.medvram: + lowvram.send_everything_to_cpu() + else: + sd_model.to(devices.cpu) + + sd_hijack.model_hijack.undo_hijack(sd_model) + + load_vae(sd_model, vae_file) + + sd_hijack.model_hijack.hijack(sd_model) + script_callbacks.model_loaded_callback(sd_model) + + if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram: + sd_model.to(devices.device) + + print(f"VAE Weights loaded.") + return sd_model diff --git a/webui.py b/webui.py index 7cb4691b..034777a2 100644 --- a/webui.py +++ b/webui.py @@ -81,9 +81,7 @@ def initialize(): modules.sd_vae.refresh_vae_list() modules.sd_models.load_model() shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: modules.sd_models.reload_model_weights())) - # I don't know what needs to be done to only reload VAE, with all those hijacks callbacks, and lowvram, - # so for now this reloads the whole model too - shared.opts.onchange("sd_vae", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(force=True)), call=False) + shared.opts.onchange("sd_vae", 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) -- cgit v1.2.3 From dd2108fdac2ebf943d4ac3563a49202222b88acf Mon Sep 17 00:00:00 2001 From: Maiko Tan Date: Wed, 2 Nov 2022 15:04:35 +0800 Subject: fix: should invoke callback as well in api only mode --- modules/script_callbacks.py | 3 ++- webui.py | 2 ++ 2 files changed, 4 insertions(+), 1 deletion(-) (limited to 'webui.py') diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py index da88635b..c28e220e 100644 --- a/modules/script_callbacks.py +++ b/modules/script_callbacks.py @@ -2,6 +2,7 @@ import sys import traceback from collections import namedtuple import inspect +from typing import Optional from fastapi import FastAPI from gradio import Blocks @@ -62,7 +63,7 @@ def clear_callbacks(): callbacks_image_saved.clear() callbacks_cfg_denoiser.clear() -def app_started_callback(demo: Blocks, app: FastAPI): +def app_started_callback(demo: Optional[Blocks], app: FastAPI): for c in callbacks_app_started: try: c.callback(demo, app) diff --git a/webui.py b/webui.py index 84e5c1fd..dc4223dc 100644 --- a/webui.py +++ b/webui.py @@ -114,6 +114,8 @@ def api_only(): app.add_middleware(GZipMiddleware, minimum_size=1000) api = create_api(app) + modules.script_callbacks.app_started_callback(None, app) + api.launch(server_name="0.0.0.0" if cmd_opts.listen else "127.0.0.1", port=cmd_opts.port if cmd_opts.port else 7861) -- cgit v1.2.3