Fix logging configuration again

* Only use `tqdm.write()` if `tqdm` is active, defer to stderr
* Correct log formatter for TqdmLoggingHandler
* If `rich` is installed and `SD_WEBUI_RICH_LOG` is set, use `rich`'s formatter

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):
+        def __init__(self, fallback_handler: logging.Handler):
-            super().__init__(level)
+            super().__init__()
+            self.fallback_handler = fallback_handler
 
         def emit(self, record):
             try:
-                msg = self.format(record)
+                # If there are active tqdm progress bars,
-                tqdm.write(msg)
+                # attempt to not interfere with them.
-                self.flush()
+                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
+    TqdmLoggingHandler = None
-    # 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 not loglevel:
+        return
 
-    if TQDM_IMPORTED:
+    if logging.root.handlers:
-        loghandlers.append(TqdmLoggingHandler())
+        # Already configured, do not interfere
+        return
 
-    if loglevel:
+    if os.environ.get("SD_WEBUI_RICH_LOG"):
-        log_level = getattr(logging, loglevel.upper(), None) or logging.INFO
+        from rich.logging import RichHandler
-        logging.basicConfig(
+        handler = RichHandler()
-            level=log_level,
+    else:
-            format='%(asctime)s %(levelname)s [%(name)s] %(message)s',
+        handler = logging.StreamHandler()
-            datefmt='%Y-%m-%d %H:%M:%S',
+
-            handlers=loghandlers
+    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)