pref: optimize print logging

py/config.py

@@ -1,3 +1,5 @@
+extension_tag = "ComfyUI Model Manager"
+
 extension_uri: str = None
 model_base_paths: dict[str, list[str]] = {}
 

py/download.py

@@ -1,10 +1,8 @@
 import os
 import uuid
 import time
-import logging
 import requests
 import folder_paths
-import traceback
 from typing import Callable, Awaitable, Any, Literal, Union, Optional
 from dataclasses import dataclass
 from . import config
@@ -225,7 +223,7 @@ async def download_model(task_id: str, request):
             task_status.error = str(e)
             await utils.send_json("update_download_task", task_status)
             task_status.error = None
-            logging.error(str(e))
+            utils.print_error(str(e))
 
     try:
         status = download_thread_pool.submit(download_task, task_id)
@@ -238,7 +236,7 @@ async def download_model(task_id: str, request):
         task_status.error = str(e)
         await utils.send_json("update_download_task", task_status)
         task_status.error = None
-        logging.error(traceback.format_exc())
+        utils.print_error(str(e))
 
 
 async def download_model_file(

py/thread.py

@@ -1,7 +1,8 @@
 import asyncio
 import threading
 import queue
-import logging
+
+from . import utils
 
 
 class DownloadThreadPool:
@@ -50,7 +51,7 @@ class DownloadThreadPool:
                 with self._lock:
                     self.running_tasks.remove(task_id)
             except Exception as e:
-                logging.error(f"worker run error: {str(e)}")
+                utils.print_error(f"worker run error: {str(e)}")
 
         with self._lock:
             self.workers_count -= 1

py/utils.py

@@ -5,6 +5,7 @@ import shutil
 import tarfile
 import logging
 import requests
+import traceback
 import configparser
 
 import comfy.utils
@@ -15,6 +16,15 @@ from typing import Any
 from . import config
 
 
+def print_info(msg, *args, **kwargs):
+    logging.info(f"[{config.extension_tag}] {msg}", *args, **kwargs)
+
+
+def print_error(msg, *args, **kwargs):
+    logging.error(f"[{config.extension_tag}] {msg}", *args, **kwargs)
+    logging.debug(traceback.format_exc())
+
+
 def normalize_path(path: str):
     normpath = os.path.normpath(path)
     return normpath.replace(os.path.sep, "/")
@@ -52,8 +62,8 @@ def download_web_distribution(version: str):
         return
 
     try:
-        logging.info(f"current version {version}, web version {web_version}")
-        logging.info("Downloading web distribution...")
+        print_info(f"current version {version}, web version {web_version}")
+        print_info("Downloading web distribution...")
         download_url = f"https://github.com/hayden-fr/ComfyUI-Model-Manager/releases/download/v{version}/dist.tar.gz"
         response = requests.get(download_url, stream=True)
         response.raise_for_status()
@@ -66,7 +76,7 @@ def download_web_distribution(version: str):
         if os.path.exists(web_path):
             shutil.rmtree(web_path)
 
-        logging.info("Extracting web distribution...")
+        print_info("Extracting web distribution...")
         with tarfile.open(temp_file, "r:gz") as tar:
             members = [
                 member for member in tar.getmembers() if member.name.startswith("web/")
@@ -74,13 +84,13 @@ def download_web_distribution(version: str):
             tar.extractall(path=config.extension_uri, members=members)
 
         os.remove(temp_file)
-        logging.info("Web distribution downloaded successfully.")
+        print_info("Web distribution downloaded successfully.")
     except requests.exceptions.RequestException as e:
-        logging.error(f"Failed to download web distribution: {e}")
+        print_error(f"Failed to download web distribution: {e}")
     except tarfile.TarError as e:
-        logging.error(f"Failed to extract web distribution: {e}")
+        print_error(f"Failed to extract web distribution: {e}")
     except Exception as e:
-        logging.error(f"An unexpected error occurred: {e}")
+        print_error(f"An unexpected error occurred: {e}")
 
 
 def resolve_model_base_paths():