logging_config.py 1.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758
  1. import logging
  2. import os
  3. try:
  4. from tqdm import tqdm
  5. class TqdmLoggingHandler(logging.Handler):
  6. def __init__(self, fallback_handler: logging.Handler):
  7. super().__init__()
  8. self.fallback_handler = fallback_handler
  9. def emit(self, record):
  10. try:
  11. # If there are active tqdm progress bars,
  12. # attempt to not interfere with them.
  13. if tqdm._instances:
  14. tqdm.write(self.format(record))
  15. else:
  16. self.fallback_handler.emit(record)
  17. except Exception:
  18. self.fallback_handler.emit(record)
  19. except ImportError:
  20. TqdmLoggingHandler = None
  21. def setup_logging(loglevel):
  22. if loglevel is None:
  23. loglevel = os.environ.get("SD_WEBUI_LOG_LEVEL")
  24. if not loglevel:
  25. return
  26. if logging.root.handlers:
  27. # Already configured, do not interfere
  28. return
  29. formatter = logging.Formatter(
  30. '%(asctime)s %(levelname)s [%(name)s] %(message)s',
  31. '%Y-%m-%d %H:%M:%S',
  32. )
  33. if os.environ.get("SD_WEBUI_RICH_LOG"):
  34. from rich.logging import RichHandler
  35. handler = RichHandler()
  36. else:
  37. handler = logging.StreamHandler()
  38. handler.setFormatter(formatter)
  39. if TqdmLoggingHandler:
  40. handler = TqdmLoggingHandler(handler)
  41. handler.setFormatter(formatter)
  42. log_level = getattr(logging, loglevel.upper(), None) or logging.INFO
  43. logging.root.setLevel(log_level)
  44. logging.root.addHandler(handler)