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

1 files changed, 24 insertions, 13 deletions

diff --git a/modules/logging_config.py b/modules/logging_config.py
index ce831b5c..99ed2855 100644
--- a/modules/logging_config.py
+++ b/modules/logging_config.py
@@ -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=[]

         )

-