Merge pull request #4592 from awsr/errorlimiter

Implement abort system when encountering excessive errors

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)

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)
 

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
 

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}')

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):