diff --git a/modules/logging_config.py b/modules/logging_config.py index 792698756..11eee9a63 100644 --- a/modules/logging_config.py +++ b/modules/logging_config.py @@ -1,41 +1,57 @@ -import os import logging +import os try: - from tqdm.auto import tqdm + from tqdm import tqdm + class TqdmLoggingHandler(logging.Handler): - def __init__(self, level=logging.INFO): - super().__init__(level) + def __init__(self, fallback_handler: logging.Handler): + super().__init__() + self.fallback_handler = fallback_handler def emit(self, record): try: - msg = self.format(record) - tqdm.write(msg) - self.flush() + # If there are active tqdm progress bars, + # attempt to not interfere with them. + if tqdm._instances: + tqdm.write(self.format(record)) + else: + self.fallback_handler.emit(record) except Exception: - self.handleError(record) + self.fallback_handler.emit(record) - TQDM_IMPORTED = True except ImportError: - # tqdm does not exist before first launch - # I will import once the UI finishes seting up the enviroment and reloads. - TQDM_IMPORTED = False + TqdmLoggingHandler = None + def setup_logging(loglevel): if loglevel is None: loglevel = os.environ.get("SD_WEBUI_LOG_LEVEL") - loghandlers = [] + if not loglevel: + return - if TQDM_IMPORTED: - loghandlers.append(TqdmLoggingHandler()) + if logging.root.handlers: + # Already configured, do not interfere + return - if loglevel: - log_level = getattr(logging, loglevel.upper(), None) or logging.INFO - logging.basicConfig( - level=log_level, - format='%(asctime)s %(levelname)s [%(name)s] %(message)s', - datefmt='%Y-%m-%d %H:%M:%S', - handlers=loghandlers - ) + if os.environ.get("SD_WEBUI_RICH_LOG"): + from rich.logging import RichHandler + handler = RichHandler() + else: + handler = logging.StreamHandler() + + if TqdmLoggingHandler: + handler = TqdmLoggingHandler(handler) + + formatter = logging.Formatter( + '%(asctime)s %(levelname)s [%(name)s] %(message)s', + '%Y-%m-%d %H:%M:%S', + ) + + handler.setFormatter(formatter) + + log_level = getattr(logging, loglevel.upper(), None) or logging.INFO + logging.root.setLevel(log_level) + logging.root.addHandler(handler)