Merge branch 'dev' into kingljl-patch-memory-leak

1 files changed, 26 insertions, 1 deletions

diff --git a/modules/logging_config.py b/modules/logging_config.py
index 7db23d4b..79269875 100644
--- a/modules/logging_config.py
+++ b/modules/logging_config.py
@@ -1,16 +1,41 @@
 import os

 import logging

 

+try:

+    from tqdm.auto import tqdm

+

+    class TqdmLoggingHandler(logging.Handler):

+        def __init__(self, level=logging.INFO):

+            super().__init__(level)

+

+        def emit(self, record):

+            try:

+                msg = self.format(record)

+                tqdm.write(msg)

+                self.flush()

+            except Exception:

+                self.handleError(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

 

 def setup_logging(loglevel):

     if loglevel is None:

         loglevel = os.environ.get("SD_WEBUI_LOG_LEVEL")

 

+    loghandlers = []

+

+    if TQDM_IMPORTED:

+        loghandlers.append(TqdmLoggingHandler())

+

     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

         )

-