Commit cdb60a69 authored by Lucas Daniel Velazquez M's avatar Lucas Daniel Velazquez M Committed by Your Name
Browse files

Take into account tqdm not being installed before first boot for logging

parent 236eb82c
Loading
Loading
Loading
Loading
+24 −13
Original line number Diff line number Diff line
import os
import logging

try:
    from tqdm.auto import tqdm

    class TqdmLoggingHandler(logging.Handler):
@@ -15,16 +16,26 @@ class TqdmLoggingHandler(logging.Handler):
            except Exception:
                self.handleError(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

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

    loghandlers = []

    if TQDM_IMPORTED:
        loghandlers.append(TqdmLoggingHandler())

    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=[TqdmLoggingHandler()]
            handlers=[]
        )