diff --git a/modules/errorlimiter.py b/modules/errorlimiter.py new file mode 100644 index 000000000..ca8c1f5a4 --- /dev/null +++ b/modules/errorlimiter.py @@ -0,0 +1,73 @@ +from __future__ import annotations +from contextlib import contextmanager +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Iterable + + +class ErrorLimiterTrigger(BaseException): # Use BaseException to avoid being caught by "except Exception:". + def __init__(self, name: str, *args): + super().__init__(*args) + self.name = name + + +class ErrorLimiterAbort(RuntimeError): + def __init__(self, msg: str): + super().__init__(msg) + + +class ErrorLimiter: + _store: dict[str, int] = {} + + @classmethod + def start(cls, name: str, limit: int = 5): + cls._store[name] = limit + + @classmethod + def notify(cls, name: str | Iterable[str]): # Can be manually triggered if execution is spread across multiple files + if isinstance(name, str): + name = (name,) + for key in name: + if key in cls._store.keys(): + cls._store[key] = cls._store[key] - 1 + if cls._store[key] <= 0: + raise ErrorLimiterTrigger(key) + + @classmethod + def end(cls, name: str): + cls._store.pop(name) + + +@contextmanager +def limit_errors(name: str, limit: int = 5): + """Limiter for aborting execution after being triggered a specified number of times (default 5). + + >>> with limit_errors("identifier", limit=5) as elimit: + >>> while do_thing(): + >>> if (something_bad): + >>> print("Something bad happened") + >>> elimit() # In this example, raises ErrorLimiterAbort on the 5th call + >>> try: + >>> something_broken() + >>> except Exception: + >>> print("Encountered an exception") + >>> elimit() # Count is shared across all calls + + Args: + name (str): Identifier. + limit (int, optional): Abort after `limit` number of triggers. Defaults to 5. + + Raises: + ErrorLimiterAbort: Subclass of RuntimeException. + + Yields: + Callable: Notification function to indicate that an error occurred. + """ + try: + ErrorLimiter.start(name, limit) + yield lambda: ErrorLimiter.notify(name) + except ErrorLimiterTrigger as e: + raise ErrorLimiterAbort(f"HALTING. Too many errors during '{e.name}'") from None + finally: + ErrorLimiter.end(name) diff --git a/modules/errors.py b/modules/errors.py index 81cfe9379..29566d597 100644 --- a/modules/errors.py +++ b/modules/errors.py @@ -1,6 +1,7 @@ import logging import warnings from installer import get_log, get_console, setup_logging, install_traceback +from modules.errorlimiter import ErrorLimiterAbort log = get_log() @@ -17,6 +18,8 @@ def install(suppress=[]): def display(e: Exception, task: str, suppress=[]): log.error(f"{task or 'error'}: {type(e).__name__}") + if isinstance(e, ErrorLimiterAbort): + return console = get_console() console.print_exception(show_locals=False, max_frames=16, extra_lines=1, suppress=suppress, theme="ansi_dark", word_wrap=False, width=console.width) diff --git a/modules/lora/lora_apply.py b/modules/lora/lora_apply.py index 06b896349..e79306c9f 100644 --- a/modules/lora/lora_apply.py +++ b/modules/lora/lora_apply.py @@ -3,6 +3,7 @@ import re import time import torch import diffusers.models.lora +from modules.errorlimiter import ErrorLimiter from modules.lora import lora_common as l from modules import shared, devices, errors, model_quant @@ -141,6 +142,7 @@ def network_calc_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn. if l.debug: errors.display(e, 'LoRA') raise RuntimeError('LoRA apply weight') from e + ErrorLimiter.notify(("network_activate", "network_deactivate")) continue return batch_updown, batch_ex_bias diff --git a/modules/lora/networks.py b/modules/lora/networks.py index 6d37fd656..69df992cc 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -1,6 +1,7 @@ from contextlib import nullcontext import time import rich.progress as rp +from modules.errorlimiter import limit_errors from modules.lora import lora_common as l from modules.lora.lora_apply import network_apply_weights, network_apply_direct, network_backup_weights, network_calc_weights from modules import shared, devices, sd_models @@ -12,61 +13,62 @@ default_components = ['text_encoder', 'text_encoder_2', 'text_encoder_3', 'text_ def network_activate(include=[], exclude=[]): t0 = time.time() - sd_model = getattr(shared.sd_model, "pipe", shared.sd_model) - if shared.opts.diffusers_offload_mode == "sequential": - sd_models.disable_offload(sd_model) - sd_models.move_model(sd_model, device=devices.cpu) - device = None - modules = {} - components = include if len(include) > 0 else default_components - components = [x for x in components if x not in exclude] - active_components = [] - for name in components: - component = getattr(sd_model, name, None) - if component is not None and hasattr(component, 'named_modules'): - active_components.append(name) - modules[name] = list(component.named_modules()) - total = sum(len(x) for x in modules.values()) - if len(l.loaded_networks) > 0: - pbar = rp.Progress(rp.TextColumn('[cyan]Network: type=LoRA action=activate'), rp.BarColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=shared.console) - task = pbar.add_task(description='' , total=total) - else: - task = None - pbar = nullcontext() - applied_weight = 0 - applied_bias = 0 - with devices.inference_context(), pbar: - wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in l.loaded_networks) if len(l.loaded_networks) > 0 else () - applied_layers.clear() - backup_size = 0 - for component in modules.keys(): - device = getattr(sd_model, component, None).device - for _, module in modules[component]: - network_layer_name = getattr(module, 'network_layer_name', None) - current_names = getattr(module, "network_current_names", ()) - if getattr(module, 'weight', None) is None or shared.state.interrupted or (network_layer_name is None) or (current_names == wanted_names): + with limit_errors("network_activate"): + sd_model = getattr(shared.sd_model, "pipe", shared.sd_model) + if shared.opts.diffusers_offload_mode == "sequential": + sd_models.disable_offload(sd_model) + sd_models.move_model(sd_model, device=devices.cpu) + device = None + modules = {} + components = include if len(include) > 0 else default_components + components = [x for x in components if x not in exclude] + active_components = [] + for name in components: + component = getattr(sd_model, name, None) + if component is not None and hasattr(component, 'named_modules'): + active_components.append(name) + modules[name] = list(component.named_modules()) + total = sum(len(x) for x in modules.values()) + if len(l.loaded_networks) > 0: + pbar = rp.Progress(rp.TextColumn('[cyan]Network: type=LoRA action=activate'), rp.BarColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=shared.console) + task = pbar.add_task(description='' , total=total) + else: + task = None + pbar = nullcontext() + applied_weight = 0 + applied_bias = 0 + with devices.inference_context(), pbar: + wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in l.loaded_networks) if len(l.loaded_networks) > 0 else () + applied_layers.clear() + backup_size = 0 + for component in modules.keys(): + device = getattr(sd_model, component, None).device + for _, module in modules[component]: + network_layer_name = getattr(module, 'network_layer_name', None) + current_names = getattr(module, "network_current_names", ()) + if getattr(module, 'weight', None) is None or shared.state.interrupted or (network_layer_name is None) or (current_names == wanted_names): + if task is not None: + pbar.update(task, advance=1) + continue + backup_size += network_backup_weights(module, network_layer_name, wanted_names) + batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name) + if shared.opts.lora_fuse_native: + network_apply_direct(module, batch_updown, batch_ex_bias, device=device) + else: + network_apply_weights(module, batch_updown, batch_ex_bias, device=device) + if batch_updown is not None or batch_ex_bias is not None: + applied_layers.append(network_layer_name) + applied_weight += 1 if batch_updown is not None else 0 + applied_bias += 1 if batch_ex_bias is not None else 0 + batch_updown, batch_ex_bias = None, None + del batch_updown, batch_ex_bias + module.network_current_names = wanted_names if task is not None: - pbar.update(task, advance=1) - continue - backup_size += network_backup_weights(module, network_layer_name, wanted_names) - batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name) - if shared.opts.lora_fuse_native: - network_apply_direct(module, batch_updown, batch_ex_bias, device=device) - else: - network_apply_weights(module, batch_updown, batch_ex_bias, device=device) - if batch_updown is not None or batch_ex_bias is not None: - applied_layers.append(network_layer_name) - applied_weight += 1 if batch_updown is not None else 0 - applied_bias += 1 if batch_ex_bias is not None else 0 - batch_updown, batch_ex_bias = None, None - del batch_updown, batch_ex_bias - module.network_current_names = wanted_names - if task is not None: - bs = round(backup_size/1024/1024/1024, 2) if backup_size > 0 else None - pbar.update(task, advance=1, description=f'networks={len(l.loaded_networks)} modules={active_components} layers={total} weights={applied_weight} bias={applied_bias} backup={bs} device={device}') + bs = round(backup_size/1024/1024/1024, 2) if backup_size > 0 else None + pbar.update(task, advance=1, description=f'networks={len(l.loaded_networks)} modules={active_components} layers={total} weights={applied_weight} bias={applied_bias} backup={bs} device={device}') - if task is not None and len(applied_layers) == 0: - pbar.remove_task(task) # hide progress bar for no action + if task is not None and len(applied_layers) == 0: + pbar.remove_task(task) # hide progress bar for no action l.timer.activate += time.time() - t0 if l.debug and len(l.loaded_networks) > 0: shared.log.debug(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} modules={active_components} layers={total} weights={applied_weight} bias={applied_bias} backup={round(backup_size/1024/1024/1024, 2)} fuse={shared.opts.lora_fuse_native}:{shared.opts.lora_fuse_diffusers} device={device} time={l.timer.summary}') @@ -81,49 +83,49 @@ def network_deactivate(include=[], exclude=[]): if len(l.previously_loaded_networks) == 0: return t0 = time.time() - sd_model = getattr(shared.sd_model, "pipe", shared.sd_model) - if shared.opts.diffusers_offload_mode == "sequential": - sd_models.disable_offload(sd_model) - sd_models.move_model(sd_model, device=devices.cpu) - modules = {} + with limit_errors("network_deactivate"): + sd_model = getattr(shared.sd_model, "pipe", shared.sd_model) + if shared.opts.diffusers_offload_mode == "sequential": + sd_models.disable_offload(sd_model) + sd_models.move_model(sd_model, device=devices.cpu) + modules = {} - components = include if len(include) > 0 else ['text_encoder', 'text_encoder_2', 'text_encoder_3', 'unet', 'transformer'] - components = [x for x in components if x not in exclude] - active_components = [] - for name in components: - component = getattr(sd_model, name, None) - if component is not None and hasattr(component, 'named_modules'): - modules[name] = list(component.named_modules()) - active_components.append(name) - total = sum(len(x) for x in modules.values()) - if len(l.previously_loaded_networks) > 0 and l.debug: - pbar = rp.Progress(rp.TextColumn('[cyan]Network: type=LoRA action=deactivate'), rp.BarColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=shared.console) - task = pbar.add_task(description='', total=total) - else: - task = None - pbar = nullcontext() - with devices.inference_context(), pbar: - applied_layers.clear() - for component in modules.keys(): - device = getattr(sd_model, component, None).device - for _, module in modules[component]: - network_layer_name = getattr(module, 'network_layer_name', None) - if shared.state.interrupted or network_layer_name is None: + components = include if len(include) > 0 else ['text_encoder', 'text_encoder_2', 'text_encoder_3', 'unet', 'transformer'] + components = [x for x in components if x not in exclude] + active_components = [] + for name in components: + component = getattr(sd_model, name, None) + if component is not None and hasattr(component, 'named_modules'): + modules[name] = list(component.named_modules()) + active_components.append(name) + total = sum(len(x) for x in modules.values()) + if len(l.previously_loaded_networks) > 0 and l.debug: + pbar = rp.Progress(rp.TextColumn('[cyan]Network: type=LoRA action=deactivate'), rp.BarColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=shared.console) + task = pbar.add_task(description='', total=total) + else: + task = None + pbar = nullcontext() + with devices.inference_context(), pbar: + applied_layers.clear() + for component in modules.keys(): + device = getattr(sd_model, component, None).device + for _, module in modules[component]: + network_layer_name = getattr(module, 'network_layer_name', None) + if shared.state.interrupted or network_layer_name is None: + if task is not None: + pbar.update(task, advance=1) + continue + batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, use_previous=True) + if shared.opts.lora_fuse_native: + network_apply_direct(module, batch_updown, batch_ex_bias, device=device, deactivate=True) + else: + network_apply_weights(module, batch_updown, batch_ex_bias, device=device, deactivate=True) + if batch_updown is not None or batch_ex_bias is not None: + applied_layers.append(network_layer_name) + del batch_updown, batch_ex_bias + module.network_current_names = () if task is not None: - pbar.update(task, advance=1) - continue - batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, use_previous=True) - if shared.opts.lora_fuse_native: - network_apply_direct(module, batch_updown, batch_ex_bias, device=device, deactivate=True) - else: - network_apply_weights(module, batch_updown, batch_ex_bias, device=device, deactivate=True) - if batch_updown is not None or batch_ex_bias is not None: - applied_layers.append(network_layer_name) - del batch_updown, batch_ex_bias - module.network_current_names = () - if task is not None: - pbar.update(task, advance=1, description=f'networks={len(l.previously_loaded_networks)} modules={active_components} layers={total} unapply={len(applied_layers)}') - + pbar.update(task, advance=1, description=f'networks={len(l.previously_loaded_networks)} modules={active_components} layers={total} unapply={len(applied_layers)}') l.timer.deactivate = time.time() - t0 if l.debug and len(l.previously_loaded_networks) > 0: shared.log.debug(f'Network deactivate: type=LoRA networks={[n.name for n in l.previously_loaded_networks]} modules={active_components} layers={total} apply={len(applied_layers)} fuse={shared.opts.lora_fuse_native}:{shared.opts.lora_fuse_diffusers} time={l.timer.summary}') diff --git a/modules/textual_inversion.py b/modules/textual_inversion.py index 4d7b76a77..064d7d214 100644 --- a/modules/textual_inversion.py +++ b/modules/textual_inversion.py @@ -3,6 +3,7 @@ import os import time import torch import safetensors.torch +from modules.errorlimiter import limit_errors from modules import shared, devices, errors from modules.files_cache import directory_files, directory_mtime, extension_filter @@ -258,47 +259,50 @@ class EmbeddingDatabase: File names take precidence over bundled embeddings passed as a dict. Bundled embeddings are automatically set to overwrite previous embeddings. """ - overwrite = bool(data) - if not shared.sd_loaded: - return - if not shared.opts.diffusers_enable_embed: - return - embeddings, skipped = open_embeddings(filename) or convert_bundled(data) - for skip in skipped: - self.skipped_embeddings[skip.name] = skipped - if not embeddings: - return - text_encoders, tokenizers, hiddensizes = get_text_encoders() - if not all([text_encoders, tokenizers, hiddensizes]): - return - for embedding in embeddings: - try: - embedding.vector_sizes = [v.shape[-1] for v in embedding.vec] - if shared.opts.diffusers_convert_embed and 768 in hiddensizes and 1280 in hiddensizes and 1280 not in embedding.vector_sizes and 768 in embedding.vector_sizes: - embedding.vec.append(convert_embedding(embedding.vec[embedding.vector_sizes.index(768)], text_encoders[hiddensizes.index(768)], text_encoders[hiddensizes.index(1280)])) - embedding.vector_sizes.append(1280) - if (not all(vs in hiddensizes for vs in embedding.vector_sizes) or # Skip SD2.1 in SD1.5/SDXL/SD3 vis versa - len(embedding.vector_sizes) > len(hiddensizes) or # Skip SDXL/SD3 in SD1.5 - (len(embedding.vector_sizes) < len(hiddensizes) and len(embedding.vector_sizes) != 2)): # SD3 no T5 - embedding.tokens = [] + with limit_errors("load_diffusers_embedding") as elimit: + overwrite = bool(data) + if not shared.sd_loaded: + return + if not shared.opts.diffusers_enable_embed: + return + embeddings, skipped = open_embeddings(filename) or convert_bundled(data) + for skip in skipped: + self.skipped_embeddings[skip.name] = skipped + if not embeddings: + return + text_encoders, tokenizers, hiddensizes = get_text_encoders() + if not all([text_encoders, tokenizers, hiddensizes]): + return + for embedding in embeddings: + try: + embedding.vector_sizes = [v.shape[-1] for v in embedding.vec] + if shared.opts.diffusers_convert_embed and 768 in hiddensizes and 1280 in hiddensizes and 1280 not in embedding.vector_sizes and 768 in embedding.vector_sizes: + embedding.vec.append(convert_embedding(embedding.vec[embedding.vector_sizes.index(768)], text_encoders[hiddensizes.index(768)], text_encoders[hiddensizes.index(1280)])) + embedding.vector_sizes.append(1280) + if (not all(vs in hiddensizes for vs in embedding.vector_sizes) or # Skip SD2.1 in SD1.5/SDXL/SD3 vis versa + len(embedding.vector_sizes) > len(hiddensizes) or # Skip SDXL/SD3 in SD1.5 + (len(embedding.vector_sizes) < len(hiddensizes) and len(embedding.vector_sizes) != 2)): # SD3 no T5 + embedding.tokens = [] + self.skipped_embeddings[embedding.name] = embedding + except Exception as e: + shared.log.error(f'Load embedding invalid: name="{embedding.name}" fn="{filename}" {e}') self.skipped_embeddings[embedding.name] = embedding - except Exception as e: - shared.log.error(f'Load embedding invalid: name="{embedding.name}" fn="{filename}" {e}') - self.skipped_embeddings[embedding.name] = embedding - if overwrite: - shared.log.info(f"Load bundled embeddings: {list(data.keys())}") + elimit() + if overwrite: + shared.log.info(f"Load bundled embeddings: {list(data.keys())}") + for embedding in embeddings: + if embedding.name not in self.skipped_embeddings: + deref_tokenizers(embedding.tokens, tokenizers) + insert_tokens(embeddings, tokenizers) for embedding in embeddings: if embedding.name not in self.skipped_embeddings: - deref_tokenizers(embedding.tokens, tokenizers) - insert_tokens(embeddings, tokenizers) - for embedding in embeddings: - if embedding.name not in self.skipped_embeddings: - try: - insert_vectors(embedding, tokenizers, text_encoders, hiddensizes) - self.register_embedding(embedding, shared.sd_model) - except Exception as e: - shared.log.error(f'Load embedding: name="{embedding.name}" file="{embedding.filename}" {e}') - errors.display(e, f'Load embedding: name="{embedding.name}" file="{embedding.filename}"') + try: + insert_vectors(embedding, tokenizers, text_encoders, hiddensizes) + self.register_embedding(embedding, shared.sd_model) + except Exception as e: + shared.log.error(f'Load embedding: name="{embedding.name}" file="{embedding.filename}" {e}') + errors.display(e, f'Load embedding: name="{embedding.name}" file="{embedding.filename}"') + elimit() return def load_from_dir(self, embdir):