Skip to content

Commit

Permalink
Take into account tqdm not being installed before first boot for logging
Browse files Browse the repository at this point in the history
  • Loading branch information
Luxter77 authored and Your Name committed Nov 16, 2023
1 parent 236eb82 commit cdb60a6
Showing 1 changed file with 24 additions and 13 deletions.
37 changes: 24 additions & 13 deletions modules/logging_config.py
Original file line number Diff line number Diff line change
@@ -1,30 +1,41 @@
import os
import logging

from tqdm.auto import tqdm
try:
from tqdm.auto import tqdm

class TqdmLoggingHandler(logging.Handler):
def __init__(self, level=logging.INFO):
super().__init__(level)
class TqdmLoggingHandler(logging.Handler):
def __init__(self, level=logging.INFO):
super().__init__(level)

def emit(self, record):
try:
msg = self.format(record)
tqdm.write(msg)
self.flush()
except Exception:
self.handleError(record)
def emit(self, record):
try:
msg = self.format(record)
tqdm.write(msg)
self.flush()
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=[]
)

0 comments on commit cdb60a6

Please sign in to comment.