stable-diffusion-webui

1
import logging
2
import os
3

4
try:
5
    from tqdm import tqdm
6

7

8
    class TqdmLoggingHandler(logging.Handler):
9
        def __init__(self, fallback_handler: logging.Handler):
10
            super().__init__()
11
            self.fallback_handler = fallback_handler
12

13
        def emit(self, record):
14
            try:
15
                # If there are active tqdm progress bars,
16
                # attempt to not interfere with them.
17
                if tqdm._instances:
18
                    tqdm.write(self.format(record))
19
                else:
20
                    self.fallback_handler.emit(record)
21
            except Exception:
22
                self.fallback_handler.emit(record)
23

24
except ImportError:
25
    TqdmLoggingHandler = None
26

27

28
def setup_logging(loglevel):
29
    if loglevel is None:
30
        loglevel = os.environ.get("SD_WEBUI_LOG_LEVEL")
31

32
    if not loglevel:
33
        return
34

35
    if logging.root.handlers:
36
        # Already configured, do not interfere
37
        return
38

39
    formatter = logging.Formatter(
40
        '%(asctime)s %(levelname)s [%(name)s] %(message)s',
41
        '%Y-%m-%d %H:%M:%S',
42
    )
43

44
    if os.environ.get("SD_WEBUI_RICH_LOG"):
45
        from rich.logging import RichHandler
46
        handler = RichHandler()
47
    else:
48
        handler = logging.StreamHandler()
49
        handler.setFormatter(formatter)
50

51
    if TqdmLoggingHandler:
52
        handler = TqdmLoggingHandler(handler)
53

54
    handler.setFormatter(formatter)
55

56
    log_level = getattr(logging, loglevel.upper(), None) or logging.INFO
57
    logging.root.setLevel(log_level)
58
    logging.root.addHandler(handler)
59