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
This commit is contained in:
Aarni Koskela 2024-01-04 16:13:08 +02:00
parent f903b4dda3
commit 6fa42e919f
1 changed files with 39 additions and 23 deletions

View File

@ -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)