diff --git a/modules/errorlimiter.py b/modules/errorlimiter.py index ca8c1f5a4..93c96377a 100644 --- a/modules/errorlimiter.py +++ b/modules/errorlimiter.py @@ -1,44 +1,53 @@ from __future__ import annotations + from contextlib import contextmanager -from typing import TYPE_CHECKING +from threading import Lock +from typing import TYPE_CHECKING, ClassVar if TYPE_CHECKING: from collections.abc import Iterable +_instance_id = 0 +_lock = Lock() -class ErrorLimiterTrigger(BaseException): # Use BaseException to avoid being caught by "except Exception:". + +def _make_unique(name: str): + global _instance_id + with _lock: # Guard against race conditions + new_name = f"{name}__{_instance_id}" + _instance_id += 1 + return new_name + + +class _ErrorLimiterTrigger(BaseException): # Use BaseException to avoid being caught by "except Exception:". def __init__(self, name: str, *args): super().__init__(*args) - self.name = name + self.name = name.rsplit("__", 1)[0] -class ErrorLimiterAbort(RuntimeError): - def __init__(self, msg: str): - super().__init__(msg) - - -class ErrorLimiter: - _store: dict[str, int] = {} +class _ErrorLimiter: + _store: ClassVar[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) + def notify(cls, key: str): + 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) +class ErrorLimiterAbort(RuntimeError): + def __init__(self, msg: str): + super().__init__(msg) + + @contextmanager def limit_errors(name: str, limit: int = 5): """Limiter for aborting execution after being triggered a specified number of times (default 5). @@ -64,10 +73,11 @@ def limit_errors(name: str, limit: int = 5): Yields: Callable: Notification function to indicate that an error occurred. """ + name_id = _make_unique(name) try: - ErrorLimiter.start(name, limit) - yield lambda: ErrorLimiter.notify(name) - except ErrorLimiterTrigger as e: + _ErrorLimiter.start(name_id, limit) + yield lambda: _ErrorLimiter.notify(name_id) + except _ErrorLimiterTrigger as e: raise ErrorLimiterAbort(f"HALTING. Too many errors during '{e.name}'") from None finally: - ErrorLimiter.end(name) + _ErrorLimiter.end(name_id) diff --git a/modules/lora/lora_apply.py b/modules/lora/lora_apply.py index e0ac9328c..62bc4d15d 100644 --- a/modules/lora/lora_apply.py +++ b/modules/lora/lora_apply.py @@ -1,12 +1,17 @@ +from __future__ import annotations + import re import time +from typing import TYPE_CHECKING 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 from modules.logger import log +if TYPE_CHECKING: + from collections.abc import Callable + bnb = None re_network_name = re.compile(r"(.*)\s*\([0-9a-fA-F]+\)") @@ -76,7 +81,7 @@ def network_backup_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Gr return backup_size -def network_calc_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.GroupNorm | torch.nn.LayerNorm | diffusers.models.lora.LoRACompatibleLinear | diffusers.models.lora.LoRACompatibleConv, network_layer_name: str, use_previous: bool = False): +def network_calc_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.GroupNorm | torch.nn.LayerNorm | diffusers.models.lora.LoRACompatibleLinear | diffusers.models.lora.LoRACompatibleConv, network_layer_name: str, use_previous: bool = False, *, elimit: Callable[[], None] | None = None): if shared.opts.diffusers_offload_mode == "none": try: self.to(devices.device) @@ -142,7 +147,8 @@ def network_calc_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Grou if l.debug: errors.display(e, 'LoRA') raise RuntimeError('LoRA apply weight') from e - ErrorLimiter.notify(("network_activate", "network_deactivate")) + if elimit is not None: + elimit() continue return batch_updown, batch_ex_bias diff --git a/modules/lora/networks.py b/modules/lora/networks.py index 6dee600be..ea0bc27b5 100644 --- a/modules/lora/networks.py +++ b/modules/lora/networks.py @@ -18,7 +18,7 @@ def network_activate(include=None, exclude=None): if include is None: include = [] t0 = time.time() - with limit_errors("network_activate"): + with limit_errors("network_activate") as elimit: sd_model = getattr(shared.sd_model, "pipe", shared.sd_model) if shared.opts.diffusers_offload_mode == "sequential": sd_models.disable_offload(sd_model) @@ -56,7 +56,7 @@ def network_activate(include=None, exclude=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) + batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, elimit=elimit) if shared.opts.lora_fuse_native: network_apply_direct(module, batch_updown, batch_ex_bias, device=device) else: @@ -92,7 +92,7 @@ def network_deactivate(include=None, exclude=None): if len(l.previously_loaded_networks) == 0: return t0 = time.time() - with limit_errors("network_deactivate"): + with limit_errors("network_deactivate") as elimit: sd_model = getattr(shared.sd_model, "pipe", shared.sd_model) if shared.opts.diffusers_offload_mode == "sequential": sd_models.disable_offload(sd_model) @@ -124,7 +124,7 @@ def network_deactivate(include=None, exclude=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) + batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, use_previous=True, elimit=elimit) if shared.opts.lora_fuse_native: network_apply_direct(module, batch_updown, batch_ex_bias, device=device, deactivate=True) else: