From 3366e494a1147e570d8527eea19da88edb3a1e0c Mon Sep 17 00:00:00 2001 From: AUTOMATIC <16777216c@gmail.com> Date: Mon, 22 May 2023 00:13:53 +0300 Subject: option to pad prompt/neg prompt to be same length --- modules/sd_models.py | 5 +++++ 1 file changed, 5 insertions(+) (limited to 'modules/sd_models.py') diff --git a/modules/sd_models.py b/modules/sd_models.py index b1afbaa7..91b3eb11 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -508,6 +508,11 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None): timer.record("scripts callbacks") + with devices.autocast(), torch.no_grad(): + sd_model.cond_stage_model_empty_prompt = sd_model.cond_stage_model([""]) + + timer.record("calculate empty prompt") + print(f"Model loaded in {timer.summary()}.") return sd_model -- cgit v1.2.3 From 3829afec365b748e330da33b00a0e363f8c8ab71 Mon Sep 17 00:00:00 2001 From: linkoid <36754150+linkoid@users.noreply.github.com> Date: Fri, 26 May 2023 15:08:53 -0400 Subject: Remove exit() from select_checkpoint() Raising a FileNotFoundError instead. --- modules/sd_models.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) (limited to 'modules/sd_models.py') diff --git a/modules/sd_models.py b/modules/sd_models.py index 91b3eb11..1871cc97 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -164,6 +164,7 @@ def model_hash(filename): def select_checkpoint(): + """Raises `FileNotFoundError` if no checkpoints are found.""" model_checkpoint = shared.opts.sd_model_checkpoint checkpoint_info = checkpoint_alisases.get(model_checkpoint, None) @@ -171,14 +172,14 @@ def select_checkpoint(): return checkpoint_info if len(checkpoints_list) == 0: - print("No checkpoints found. When searching for checkpoints, looked at:", file=sys.stderr) + error_message = "No checkpoints found. When searching for checkpoints, looked at:" if shared.cmd_opts.ckpt is not None: - print(f" - file {os.path.abspath(shared.cmd_opts.ckpt)}", file=sys.stderr) - print(f" - directory {model_path}", file=sys.stderr) + error_message += f"\n - file {os.path.abspath(shared.cmd_opts.ckpt)}" + error_message += f"\n - directory {model_path}" if shared.cmd_opts.ckpt_dir is not None: - print(f" - directory {os.path.abspath(shared.cmd_opts.ckpt_dir)}", file=sys.stderr) - print("Can't run without a checkpoint. Find and place a .ckpt or .safetensors file into any of those locations. The program will exit.", file=sys.stderr) - exit(1) + error_message += f"\n - directory {os.path.abspath(shared.cmd_opts.ckpt_dir)}" + error_message += "Can't run without a checkpoint. Find and place a .ckpt or .safetensors file into any of those locations." + raise FileNotFoundError(error_message) checkpoint_info = next(iter(checkpoints_list.values())) if model_checkpoint is not None: -- cgit v1.2.3 From 1f0fdede176989f151da6b97bd9a140b7f0af6e5 Mon Sep 17 00:00:00 2001 From: linkoid <36754150+linkoid@users.noreply.github.com> Date: Fri, 26 May 2023 15:15:59 -0400 Subject: Show full traceback in get_sd_model() to reveal if an error is caused by an extension --- modules/errors.py | 8 ++++++-- modules/sd_models.py | 2 +- 2 files changed, 7 insertions(+), 3 deletions(-) (limited to 'modules/sd_models.py') diff --git a/modules/errors.py b/modules/errors.py index f6b80dbb..da4694f8 100644 --- a/modules/errors.py +++ b/modules/errors.py @@ -12,9 +12,13 @@ def print_error_explanation(message): print('=' * max_len, file=sys.stderr) -def display(e: Exception, task): +def display(e: Exception, task, *, full_traceback=False): print(f"{task or 'error'}: {type(e).__name__}", file=sys.stderr) - print(traceback.format_exc(), file=sys.stderr) + te = traceback.TracebackException.from_exception(e) + if full_traceback: + # include frames leading up to the try-catch block + te.stack = traceback.StackSummary(traceback.extract_stack()[:-2] + te.stack) + print(*te.format(), sep="", file=sys.stderr) message = str(e) if "copying a param with shape torch.Size([640, 1024]) from checkpoint, the shape in current model is torch.Size([640, 768])" in message: diff --git a/modules/sd_models.py b/modules/sd_models.py index 1871cc97..3e7fc7e3 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -424,7 +424,7 @@ class SdModelData: try: load_model() except Exception as e: - errors.display(e, "loading stable diffusion model") + errors.display(e, "loading stable diffusion model", full_traceback=True) print("", file=sys.stderr) print("Stable diffusion model failed to load", file=sys.stderr) self.sd_model = None -- cgit v1.2.3 From 339b5315700a469f4a9f0d5afc08ca2aca60c579 Mon Sep 17 00:00:00 2001 From: AUTOMATIC <16777216c@gmail.com> Date: Sat, 27 May 2023 15:47:33 +0300 Subject: custom unet support --- modules/processing.py | 4 +- modules/script_callbacks.py | 20 ++++++++++ modules/sd_hijack.py | 20 +++++++--- modules/sd_models.py | 4 +- modules/sd_unet.py | 92 +++++++++++++++++++++++++++++++++++++++++++++ modules/shared.py | 1 + modules/shared_items.py | 11 ++++++ webui.py | 4 ++ 8 files changed, 148 insertions(+), 8 deletions(-) create mode 100644 modules/sd_unet.py (limited to 'modules/sd_models.py') diff --git a/modules/processing.py b/modules/processing.py index 29a3743f..b75f2515 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -13,7 +13,7 @@ from skimage import exposure from typing import Any, Dict, List import modules.sd_hijack -from modules import devices, prompt_parser, masking, sd_samplers, lowvram, generation_parameters_copypaste, extra_networks, sd_vae_approx, scripts, sd_samplers_common +from modules import devices, prompt_parser, masking, sd_samplers, lowvram, generation_parameters_copypaste, extra_networks, sd_vae_approx, scripts, sd_samplers_common, sd_unet from modules.sd_hijack import model_hijack from modules.shared import opts, cmd_opts, state import modules.shared as shared @@ -674,6 +674,8 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: if shared.opts.live_previews_enable and opts.show_progress_type == "Approx NN": sd_vae_approx.model() + sd_unet.apply_unet() + if state.job_count == -1: state.job_count = p.n_iter diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py index 40f388a5..d2728e12 100644 --- a/modules/script_callbacks.py +++ b/modules/script_callbacks.py @@ -111,6 +111,7 @@ callback_map = dict( callbacks_before_ui=[], callbacks_on_reload=[], callbacks_list_optimizers=[], + callbacks_list_unets=[], ) @@ -271,6 +272,18 @@ def list_optimizers_callback(): return res +def list_unets_callback(): + res = [] + + for c in callback_map['callbacks_list_unets']: + try: + c.callback(res) + except Exception: + report_exception(c, 'list_unets') + + return res + + def add_callback(callbacks, fun): stack = [x for x in inspect.stack() if x.filename != __file__] filename = stack[0].filename if len(stack) > 0 else 'unknown file' @@ -430,3 +443,10 @@ def on_list_optimizers(callback): to it.""" add_callback(callback_map['callbacks_list_optimizers'], callback) + + +def on_list_unets(callback): + """register a function to be called when UI is making a list of alternative options for unet. + The function will be called with one argument, a list, and shall add objects of type modules.sd_unet.SdUnetOption to it.""" + + add_callback(callback_map['callbacks_list_unets'], callback) diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py index f93df0a6..487dfd60 100644 --- a/modules/sd_hijack.py +++ b/modules/sd_hijack.py @@ -3,7 +3,7 @@ from torch.nn.functional import silu from types import MethodType import modules.textual_inversion.textual_inversion -from modules import devices, sd_hijack_optimizations, shared, script_callbacks, errors +from modules import devices, sd_hijack_optimizations, shared, script_callbacks, errors, sd_unet from modules.hypernetworks import hypernetwork from modules.shared import cmd_opts from modules import sd_hijack_clip, sd_hijack_open_clip, sd_hijack_unet, sd_hijack_xlmr, xlmr @@ -43,7 +43,7 @@ def list_optimizers(): optimizers.extend(new_optimizers) -def apply_optimizations(): +def apply_optimizations(option=None): global current_optimizer undo_optimizations() @@ -60,7 +60,7 @@ def apply_optimizations(): current_optimizer.undo() current_optimizer = None - selection = shared.opts.cross_attention_optimization + selection = option or shared.opts.cross_attention_optimization if selection == "Automatic" and len(optimizers) > 0: matching_optimizer = next(iter([x for x in optimizers if x.cmd_opt and getattr(shared.cmd_opts, x.cmd_opt, False)]), optimizers[0]) else: @@ -72,12 +72,13 @@ def apply_optimizations(): matching_optimizer = optimizers[0] if matching_optimizer is not None: - print(f"Applying optimization: {matching_optimizer.name}... ", end='') + print(f"Applying attention optimization: {matching_optimizer.name}... ", end='') matching_optimizer.apply() print("done.") current_optimizer = matching_optimizer return current_optimizer.name else: + print("Disabling attention optimization") return '' @@ -155,9 +156,9 @@ class StableDiffusionModelHijack: def __init__(self): self.embedding_db.add_embedding_dir(cmd_opts.embeddings_dir) - def apply_optimizations(self): + def apply_optimizations(self, option=None): try: - self.optimization_method = apply_optimizations() + self.optimization_method = apply_optimizations(option) except Exception as e: errors.display(e, "applying cross attention optimization") undo_optimizations() @@ -194,6 +195,11 @@ class StableDiffusionModelHijack: self.layers = flatten(m) + if not hasattr(ldm.modules.diffusionmodules.openaimodel, 'copy_of_UNetModel_forward_for_webui'): + ldm.modules.diffusionmodules.openaimodel.copy_of_UNetModel_forward_for_webui = ldm.modules.diffusionmodules.openaimodel.UNetModel.forward + + ldm.modules.diffusionmodules.openaimodel.UNetModel.forward = sd_unet.UNetModel_forward + def undo_hijack(self, m): if type(m.cond_stage_model) == xlmr.BertSeriesModelWithTransformation: m.cond_stage_model = m.cond_stage_model.wrapped @@ -215,6 +221,8 @@ class StableDiffusionModelHijack: self.layers = None self.clip = None + ldm.modules.diffusionmodules.openaimodel.UNetModel.forward = ldm.modules.diffusionmodules.openaimodel.copy_of_UNetModel_forward_for_webui + def apply_circular(self, enable): if self.circular_enabled == enable: return diff --git a/modules/sd_models.py b/modules/sd_models.py index 91b3eb11..835bc016 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -14,7 +14,7 @@ import ldm.modules.midas as midas from ldm.util import instantiate_from_config -from modules import paths, shared, modelloader, devices, script_callbacks, sd_vae, sd_disable_initialization, errors, hashes, sd_models_config +from modules import paths, shared, modelloader, devices, script_callbacks, sd_vae, sd_disable_initialization, errors, hashes, sd_models_config, sd_unet from modules.sd_hijack_inpainting import do_inpainting_hijack from modules.timer import Timer import tomesd @@ -532,6 +532,8 @@ def reload_model_weights(sd_model=None, info=None): if sd_model.sd_model_checkpoint == checkpoint_info.filename: return + sd_unet.apply_unet("None") + if shared.cmd_opts.lowvram or shared.cmd_opts.medvram: lowvram.send_everything_to_cpu() else: diff --git a/modules/sd_unet.py b/modules/sd_unet.py new file mode 100644 index 00000000..6d708ad2 --- /dev/null +++ b/modules/sd_unet.py @@ -0,0 +1,92 @@ +import torch.nn +import ldm.modules.diffusionmodules.openaimodel + +from modules import script_callbacks, shared, devices + +unet_options = [] +current_unet_option = None +current_unet = None + + +def list_unets(): + new_unets = script_callbacks.list_unets_callback() + + unet_options.clear() + unet_options.extend(new_unets) + + +def get_unet_option(option=None): + option = option or shared.opts.sd_unet + + if option == "None": + return None + + if option == "Automatic": + name = shared.sd_model.sd_checkpoint_info.model_name + + options = [x for x in unet_options if x.model_name == name] + + option = options[0].label if options else "None" + + return next(iter([x for x in unet_options if x.label == option]), None) + + +def apply_unet(option=None): + global current_unet_option + global current_unet + + new_option = get_unet_option(option) + if new_option == current_unet_option: + return + + if current_unet is not None: + print(f"Dectivating unet: {current_unet.option.label}") + current_unet.deactivate() + + current_unet_option = new_option + if current_unet_option is None: + current_unet = None + + if not (shared.cmd_opts.lowvram or shared.cmd_opts.medvram): + shared.sd_model.model.diffusion_model.to(devices.device) + + return + + shared.sd_model.model.diffusion_model.to(devices.cpu) + devices.torch_gc() + + current_unet = current_unet_option.create_unet() + current_unet.option = current_unet_option + print(f"Activating unet: {current_unet.option.label}") + current_unet.activate() + + +class SdUnetOption: + model_name = None + """name of related checkpoint - this option will be selected automatically for unet if the name of checkpoint matches this""" + + label = None + """name of the unet in UI""" + + def create_unet(self): + """returns SdUnet object to be used as a Unet instead of built-in unet when making pictures""" + raise NotImplementedError() + + +class SdUnet(torch.nn.Module): + def forward(self, x, timesteps, context, *args, **kwargs): + raise NotImplementedError() + + def activate(self): + pass + + def deactivate(self): + pass + + +def UNetModel_forward(self, x, timesteps=None, context=None, *args, **kwargs): + if current_unet is not None: + return current_unet.forward(x, timesteps, context, *args, **kwargs) + + return ldm.modules.diffusionmodules.openaimodel.copy_of_UNetModel_forward_for_webui(self, x, timesteps, context, *args, **kwargs) + diff --git a/modules/shared.py b/modules/shared.py index 0897f937..a5e7824a 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -403,6 +403,7 @@ options_templates.update(options_section(('sd', "Stable Diffusion"), { "sd_vae_checkpoint_cache": OptionInfo(0, "VAE Checkpoints to cache in RAM", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}), "sd_vae": OptionInfo("Automatic", "SD VAE", gr.Dropdown, lambda: {"choices": shared_items.sd_vae_items()}, refresh=shared_items.refresh_vae_list).info("choose VAE model: Automatic = use one with same filename as checkpoint; None = use VAE from checkpoint"), "sd_vae_as_default": OptionInfo(True, "Ignore selected VAE for stable diffusion checkpoints that have their own .vae.pt next to them"), + "sd_unet": OptionInfo("Automatic", "SD Unet", gr.Dropdown, lambda: {"choices": shared_items.sd_unet_items()}, refresh=shared_items.refresh_unet_list).info("choose Unet model: Automatic = use one with same filename as checkpoint; None = use Unet from checkpoint"), "inpainting_mask_weight": OptionInfo(1.0, "Inpainting conditioning mask strength", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}), "initial_noise_multiplier": OptionInfo(1.0, "Noise multiplier for img2img", gr.Slider, {"minimum": 0.5, "maximum": 1.5, "step": 0.01}), "img2img_color_correction": OptionInfo(False, "Apply color correction to img2img results to match original colors."), diff --git a/modules/shared_items.py b/modules/shared_items.py index 2a8713c8..7f306a06 100644 --- a/modules/shared_items.py +++ b/modules/shared_items.py @@ -29,3 +29,14 @@ def cross_attention_optimizations(): return ["Automatic"] + [x.title() for x in modules.sd_hijack.optimizers] + ["None"] +def sd_unet_items(): + import modules.sd_unet + + return ["Automatic"] + [x.label for x in modules.sd_unet.unet_options] + ["None"] + + +def refresh_unet_list(): + import modules.sd_unet + + modules.sd_unet.list_unets() + diff --git a/webui.py b/webui.py index f9210f41..1e3ff061 100644 --- a/webui.py +++ b/webui.py @@ -58,6 +58,7 @@ import modules.sd_hijack import modules.sd_hijack_optimizations import modules.sd_models import modules.sd_vae +import modules.sd_unet import modules.txt2img import modules.script_callbacks import modules.textual_inversion.textual_inversion @@ -291,6 +292,9 @@ def initialize_rest(*, reload_script_modules=False): modules.sd_hijack.list_optimizers() startup_timer.record("scripts list_optimizers") + modules.sd_unet.list_unets() + startup_timer.record("scripts list_unets") + def load_model(): """ Accesses shared.sd_model property to load model. -- cgit v1.2.3 From f1533de982350af06bc9cbaa436b3e4dfdef4eb8 Mon Sep 17 00:00:00 2001 From: AUTOMATIC <16777216c@gmail.com> Date: Thu, 1 Jun 2023 07:28:13 +0300 Subject: assign devices.dtype early because it's needed before the model is loaded --- modules/sd_models.py | 2 -- modules/shared.py | 4 ++++ 2 files changed, 4 insertions(+), 2 deletions(-) (limited to 'modules/sd_models.py') diff --git a/modules/sd_models.py b/modules/sd_models.py index 232eb9c4..918f6fd6 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -314,8 +314,6 @@ def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer timer.record("apply half()") - 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 devices.dtype_unet = model.model.diffusion_model.dtype devices.unet_needs_upcast = shared.cmd_opts.upcast_sampling and devices.dtype == torch.float16 and devices.dtype_unet == torch.float16 diff --git a/modules/shared.py b/modules/shared.py index 8ba1c569..7025a754 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -6,6 +6,7 @@ import threading import time import gradio as gr +import torch import tqdm import modules.interrogate @@ -63,6 +64,9 @@ cmd_opts.disable_extension_access = (cmd_opts.share or cmd_opts.listen or cmd_op devices.device, devices.device_interrogate, devices.device_gfpgan, devices.device_esrgan, devices.device_codeformer = \ (devices.cpu if any(y in cmd_opts.use_cpu for y in [x, 'all']) else devices.get_optimal_device() for x in ['sd', 'interrogate', 'gfpgan', 'esrgan', 'codeformer']) +devices.dtype = torch.float32 if cmd_opts.no_half else torch.float16 +devices.dtype_vae = torch.float32 if cmd_opts.no_half or cmd_opts.no_half_vae else torch.float16 + device = devices.device weight_load_location = None if cmd_opts.lowram else "cpu" -- cgit v1.2.3 From 165ab44f03cc17dc3e4c35b3e5976f3d646c7ac7 Mon Sep 17 00:00:00 2001 From: Aarni Koskela Date: Mon, 29 May 2023 10:18:15 +0300 Subject: Use os.makedirs(..., exist_ok=True) --- modules/codeformer_model.py | 4 +--- modules/extensions.py | 3 +-- modules/gfpgan_model.py | 5 +---- modules/modelloader.py | 3 +-- modules/sd_models.py | 3 +-- modules/textual_inversion/autocrop.py | 3 +-- 6 files changed, 6 insertions(+), 15 deletions(-) (limited to 'modules/sd_models.py') diff --git a/modules/codeformer_model.py b/modules/codeformer_model.py index 4260b016..d974e4b8 100644 --- a/modules/codeformer_model.py +++ b/modules/codeformer_model.py @@ -20,9 +20,7 @@ codeformer = None def setup_model(dirname): - global model_path - if not os.path.exists(model_path): - os.makedirs(model_path) + os.makedirs(model_path, exist_ok=True) path = modules.paths.paths.get("CodeFormer", None) if path is None: diff --git a/modules/extensions.py b/modules/extensions.py index 8608584b..abc6e2b1 100644 --- a/modules/extensions.py +++ b/modules/extensions.py @@ -7,8 +7,7 @@ from modules.paths_internal import extensions_dir, extensions_builtin_dir, scrip extensions = [] -if not os.path.exists(extensions_dir): - os.makedirs(extensions_dir) +os.makedirs(extensions_dir, exist_ok=True) def active(): diff --git a/modules/gfpgan_model.py b/modules/gfpgan_model.py index e239a09d..6ecd295c 100644 --- a/modules/gfpgan_model.py +++ b/modules/gfpgan_model.py @@ -70,11 +70,8 @@ gfpgan_constructor = None def setup_model(dirname): - global model_path - if not os.path.exists(model_path): - os.makedirs(model_path) - try: + os.makedirs(model_path, exist_ok=True) from gfpgan import GFPGANer from facexlib import detection, parsing # noqa: F401 global user_path diff --git a/modules/modelloader.py b/modules/modelloader.py index be23071a..75f01247 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -95,8 +95,7 @@ def cleanup_models(): def move_files(src_path: str, dest_path: str, ext_filter: str = None): try: - if not os.path.exists(dest_path): - os.makedirs(dest_path) + os.makedirs(dest_path, exist_ok=True) if os.path.exists(src_path): for file in os.listdir(src_path): fullpath = os.path.join(src_path, file) diff --git a/modules/sd_models.py b/modules/sd_models.py index 918f6fd6..6ff5d17d 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -95,8 +95,7 @@ except Exception: def setup_model(): - if not os.path.exists(model_path): - os.makedirs(model_path) + os.makedirs(model_path, exist_ok=True) enable_midas_autodownload() diff --git a/modules/textual_inversion/autocrop.py b/modules/textual_inversion/autocrop.py index 75705459..1675e39a 100644 --- a/modules/textual_inversion/autocrop.py +++ b/modules/textual_inversion/autocrop.py @@ -298,8 +298,7 @@ def download_and_cache_models(dirname): download_url = 'https://github.com/opencv/opencv_zoo/blob/91fb0290f50896f38a0ab1e558b74b16bc009428/models/face_detection_yunet/face_detection_yunet_2022mar.onnx?raw=true' model_file_name = 'face_detection_yunet.onnx' - if not os.path.exists(dirname): - os.makedirs(dirname) + os.makedirs(dirname, exist_ok=True) cache_file = os.path.join(dirname, model_file_name) if not os.path.exists(cache_file): -- cgit v1.2.3