Move to modules

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,7 +1,7 @@
 import logging
 import warnings
 from installer import get_log, get_console, setup_logging, install_traceback
-from sdnext_core.errorlimiter import ErrorLimiterAbort
+from modules.errorlimiter import ErrorLimiterAbort
 
 
 log = get_log()

modules/lora/lora_apply.py

@@ -3,7 +3,7 @@ import re
 import time
 import torch
 import diffusers.models.lora
-from sdnext_core.errorlimiter import ErrorLimiter
+from modules.errorlimiter import ErrorLimiter
 from modules.lora import lora_common as l
 from modules import shared, devices, errors, model_quant
 

modules/lora/networks.py

@@ -1,7 +1,7 @@
 from contextlib import nullcontext
 import time
 import rich.progress as rp
-from sdnext_core.errorlimiter import limit_errors
+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

modules/textual_inversion.py

@@ -3,7 +3,7 @@ import os
 import time
 import torch
 import safetensors.torch
-from sdnext_core.errorlimiter import limit_errors
+from modules.errorlimiter import limit_errors
 from modules import shared, devices, errors
 from modules.files_cache import directory_files, directory_mtime, extension_filter