logging_config.py 1.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041
  1. import os
  2. import logging
  3. try:
  4. from tqdm.auto import tqdm
  5. class TqdmLoggingHandler(logging.Handler):
  6. def __init__(self, level=logging.INFO):
  7. super().__init__(level)
  8. def emit(self, record):
  9. try:
  10. msg = self.format(record)
  11. tqdm.write(msg)
  12. self.flush()
  13. except Exception:
  14. self.handleError(record)
  15. TQDM_IMPORTED = True
  16. except ImportError:
  17. # tqdm does not exist before first launch
  18. # I will import once the UI finishes seting up the enviroment and reloads.
  19. TQDM_IMPORTED = False
  20. def setup_logging(loglevel):
  21. if loglevel is None:
  22. loglevel = os.environ.get("SD_WEBUI_LOG_LEVEL")
  23. loghandlers = []
  24. if TQDM_IMPORTED:
  25. loghandlers.append(TqdmLoggingHandler())
  26. if loglevel:
  27. log_level = getattr(logging, loglevel.upper(), None) or logging.INFO
  28. logging.basicConfig(
  29. level=log_level,
  30. format='%(asctime)s %(levelname)s [%(name)s] %(message)s',
  31. datefmt='%Y-%m-%d %H:%M:%S',
  32. handlers=loghandlers
  33. )