From f5188559e5d5ce289a4c2a17868a05cd81d054b1 Mon Sep 17 00:00:00 2001 From: Tung Nguyen Date: Thu, 27 Jul 2023 02:41:46 +0700 Subject: [PATCH] chore: improve logger --- __init__.py | 10 +++++----- config.py | 13 +++++++++++-- modules/art_venture.py | 32 ++++++++++++++++---------------- modules/log.py | 28 ---------------------------- modules/logger.py | 36 ++++++++++++++++++++++++++++++++++++ modules/nodes.py | 4 ++-- modules/utils.py | 12 ++++++------ modules/workflow.py | 2 +- 8 files changed, 77 insertions(+), 60 deletions(-) delete mode 100644 modules/log.py create mode 100644 modules/logger.py diff --git a/__init__.py b/__init__.py index 0948613..b1041f4 100644 --- a/__init__.py +++ b/__init__.py @@ -6,7 +6,7 @@ from types import MethodType from typing import Callable from .config import get_ext_dir, get_comfy_dir -from .modules.log import logger as log +from .modules.logger import logger from .modules.nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS from .modules.workflow import update_checkpoints_hash from .modules.art_venture import ArtVentureRunner @@ -52,20 +52,20 @@ def is_junction(path): def install_js(): src_dir = get_ext_dir("javascript") if not os.path.exists(src_dir): - log.error("No JS") + logger.error("No JS") return dst_dir = get_web_ext_dir() if os.path.exists(dst_dir): if os.path.islink(dst_dir) or is_junction(dst_dir): - log.info("JS already linked") + logger.info("JS already linked") return elif link_js(src_dir, dst_dir): - log.info("JS linked") + logger.info("JS linked") return - log.info("Unable to make symlink, copying JS files") + logger.info("Unable to make symlink, copying JS files") shutil.copytree(src_dir, dst_dir, dirs_exist_ok=True) diff --git a/config.py b/config.py index 753e2ea..541aab5 100644 --- a/config.py +++ b/config.py @@ -5,7 +5,7 @@ from typing import Dict from server import PromptServer -from .modules.log import logger +from .modules.logger import logger comfy_dir = os.path.dirname(inspect.getfile(PromptServer)) ext_dir = os.path.dirname(os.path.realpath(__file__)) @@ -42,7 +42,7 @@ def load_config() -> Dict: "av_endpoint": "https://api.artventure.ai", "av_token": "", "runner_enabled": False, - "remove_runner_images_after_upload": False + "remove_runner_images_after_upload": False, } if not os.path.isfile(config_path): @@ -52,6 +52,15 @@ def load_config() -> Dict: with open(config_path, "r") as f: config = json.load(f) + need_update = False + for key, value in default_config.items(): + if key not in config: + config[key] = value + need_update = True + + if need_update: + write_config(config) + logger.debug(f"Loaded config {config}") return config diff --git a/modules/art_venture.py b/modules/art_venture.py index 178181b..a9c0ac6 100644 --- a/modules/art_venture.py +++ b/modules/art_venture.py @@ -9,7 +9,7 @@ from typing import Callable, Dict, List from types import MethodType from ..config import config -from .log import logger as log +from .logger import logger from .utils import get_task_from_av, upload_to_av import folder_paths @@ -105,7 +105,7 @@ class ArtVentureRunner: try: data = get_task_from_av() except Exception as e: - log.error(f"Error while getting new task {e}") + logger.error(f"Error while getting new task {e}") return (None, e) if data["has_task"] != True: @@ -113,12 +113,12 @@ class ArtVentureRunner: prompt = data.get("prompt") callback_url: str = data.get("callback_url") - log.info(f"Got new task") - log.debug(prompt) + logger.info(f"Got new task") + logger.debug(prompt) valid = validate_prompt(prompt) if not valid[0]: - log.error(f"Invalid recipe: {valid[3]}") + logger.error(f"Invalid recipe: {valid[3]}") return (callback_url, Exception("Invalid recipe")) task_id = str(uuid4()) @@ -127,14 +127,14 @@ class ArtVentureRunner: (0, task_id, prompt, {}, outputs_to_execute) ) - log.info(f"Task registered with id {task_id}") + logger.info(f"Task registered with id {task_id}") self.current_task_id = task_id self.callback_url = callback_url self.current_task_exception = None return (callback_url, None) def watching_for_new_task(self, get_task: Callable): - log.info("Watching for new task") + logger.info("Watching for new task") failed_attempts = 0 while True: @@ -145,19 +145,19 @@ class ArtVentureRunner: try: callback_url, e = get_task() if callback_url and e is not None: - log.error("Error while getting new task") - log.error(e) + logger.error("Error while getting new task") + logger.error(e) update_task_result(callback_url, False) failed_attempts += 1 else: failed_attempts = 0 except requests.exceptions.ConnectionError: - log.error("Connection error while getting new task") + logger.error("Connection error while getting new task") failed_attempts += 1 except Exception as e: - log.error("Error while getting new task") - log.error(e) - log.debug(traceback.format_exc()) + logger.error("Error while getting new task") + logger.error(e) + logger.debug(traceback.format_exc()) failed_attempts += 1 # increase sleep time based on failed attempts @@ -165,7 +165,7 @@ class ArtVentureRunner: def watching_for_new_task_threading(self): if config.get("runner_enabled", False) != True: - log.info("Runner is disabled") + logger.info("Runner is disabled") return if self.current_thread is not None and self.current_thread.is_alive(): @@ -186,7 +186,7 @@ class ArtVentureRunner: return if self.current_task_exception is not None: - log.info(f"Task {task_id} failed: {self.current_task_exception}") + logger.info(f"Task {task_id} failed: {self.current_task_exception}") update_task_result(callback_url=self.callback_url, success=False) else: images = [] @@ -200,7 +200,7 @@ class ArtVentureRunner: subfolder = image.get("subfolder", "") images.append(os.path.join(outdir, subfolder, filename)) - log.info(f"Task {task_id} finished with {len(images)} image(s)") + logger.info(f"Task {task_id} finished with {len(images)} image(s)") update_task_result(self.callback_url, True, images) if config.get("remove_runner_images_after_upload", False): for img in images: diff --git a/modules/log.py b/modules/log.py deleted file mode 100644 index 1a289f6..0000000 --- a/modules/log.py +++ /dev/null @@ -1,28 +0,0 @@ -import sys -import logging - -if logging.getLogger().hasHandlers(): - logger = logging.getLogger("sd") -else: - - class Log: - def __init__(self, level=logging.INFO) -> None: - self.level = level - - def __log(self, level, *args, **kwargs): - if (level >= self.level): - print("[ArtVenture]", *args, **kwargs) - - def info(self, *args, **kwargs): - self.__log(logging.INFO, *args, **kwargs) - - def debug(self, *args, **kwargs): - self.__log(logging.DEBUG, *args, **kwargs) - - def warning(self, *args, **kwargs): - self.__log(logging.WARNING, *args, **kwargs) - - def error(self, *args, **kwargs): - self.__log(logging.ERROR, *args, **kwargs, file=sys.stderr) - - logger = Log(level=logging.INFO) \ No newline at end of file diff --git a/modules/logger.py b/modules/logger.py new file mode 100644 index 0000000..9983e14 --- /dev/null +++ b/modules/logger.py @@ -0,0 +1,36 @@ +import sys +import copy +import logging + + +class ColoredFormatter(logging.Formatter): + COLORS = { + "DEBUG": "\033[0;36m", # CYAN + "INFO": "\033[0;32m", # GREEN + "WARNING": "\033[0;33m", # YELLOW + "ERROR": "\033[0;31m", # RED + "CRITICAL": "\033[0;37;41m", # WHITE ON RED + "RESET": "\033[0m", # RESET COLOR + } + + def format(self, record): + colored_record = copy.copy(record) + levelname = colored_record.levelname + seq = self.COLORS.get(levelname, self.COLORS["RESET"]) + colored_record.levelname = f"{seq}{levelname}{self.COLORS['RESET']}" + return super().format(colored_record) + + +# Create a new logger +logger = logging.getLogger("ArtVenture") +logger.propagate = False + +# Add handler if we don't have one. +if not logger.handlers: + handler = logging.StreamHandler(sys.stdout) + handler.setFormatter(ColoredFormatter("[%(name)s] - %(levelname)s - %(message)s")) + logger.addHandler(handler) + +# Configure logger +loglevel = logging.INFO +logger.setLevel(loglevel) diff --git a/modules/nodes.py b/modules/nodes.py index f2c8d03..ff01de1 100644 --- a/modules/nodes.py +++ b/modules/nodes.py @@ -8,7 +8,7 @@ from PIL import Image, ImageOps from PIL.PngImagePlugin import PngInfo import numpy as np -from .log import logger as log +from .logger import logger from .utils import upload_to_av import folder_paths @@ -137,7 +137,7 @@ class AVOutputUploadImage: metadata.add_text("prompt", json.dumps(prompt)) if extra_pnginfo is not None: for x in extra_pnginfo: - log.debug(f"Adding {x} to pnginfo: {extra_pnginfo[x]}") + logger.debug(f"Adding {x} to pnginfo: {extra_pnginfo[x]}") metadata.add_text(x, json.dumps(extra_pnginfo[x])) buffer = io.BytesIO() diff --git a/modules/utils.py b/modules/utils.py index e2cdc84..afa0092 100644 --- a/modules/utils.py +++ b/modules/utils.py @@ -4,7 +4,7 @@ import traceback from typing import Callable, Dict from ..config import config -from .log import logger as log +from .logger import logger def request_with_retry( @@ -19,21 +19,21 @@ def request_with_retry( return True except requests.exceptions.ConnectionError: - log.error("Connection error") + logger.error("Connection error") if retries >= max_try - 1: return False time.sleep(2) - log.info(f"Retrying {retries + 1}...") + logger.info(f"Retrying {retries + 1}...") return request_with_retry( make_request, max_try=max_try, retries=retries + 1, ) except Exception as e: - log.error("Request error") - log.error(e) - log.debug(traceback.format_exc()) + logger.error("Request error") + logger.error(e) + logger.debug(traceback.format_exc()) return False diff --git a/modules/workflow.py b/modules/workflow.py index 47d1f5a..165bc7d 100644 --- a/modules/workflow.py +++ b/modules/workflow.py @@ -10,7 +10,7 @@ from typing import Dict, List from server import PromptServer from folder_paths import models_dir, get_filename_list, get_full_path -from .log import logger +from .logger import logger from .nodes import NODE_CLASS_MAPPINGS as _NODE_CLASS_MAPPINGS node_class_mappings_loaded = False