utils/log.py

"""Console logging helpers shared by Moxie Multimedia nodes."""

from __future__ import annotations

from collections.abc import Callable, Iterator
from contextlib import contextmanager
from functools import wraps
import inspect
import logging
from typing import Any
from time import perf_counter

_BOLD_CYAN = "\033[1m\033[36m"
_RESET = "\033[0m"

__all__ = [
    "log_node_info", "log_stage_time", "synchronize_execution_device",
    "instrument_node_timing",
]


def log_node_info(node_name: str, message: str | None = None) -> None:
    """Display an informational node message in the ComfyUI console."""
    title = node_name.replace(" (EasyMedia)", "")
    suffix = f":{_RESET} {message}" if message is not None else _RESET
    print(f"{_BOLD_CYAN}[Moxie] {title}{suffix}", flush=True)  # noqa: T201


@contextmanager
def log_stage_time(
    node_name: str,
    stage: str,
    *,
    synchronize: Callable[[], None] | None = None,
) -> Iterator[None]:
    """Measure actual work, excluding upstream execution and prior GPU work."""
    if synchronize is not None:
        synchronize()
    started = perf_counter()
    try:
        yield
        if synchronize is not None:
            synchronize()
    except BaseException:
        log_node_info(node_name, f"Timing | {stage} | failed after {perf_counter() - started:.3f} s")
        raise
    else:
        log_node_info(node_name, f"Timing | {stage} | {perf_counter() - started:.3f} s")


def synchronize_execution_device() -> None:
    """Wait on the active ComfyUI device so GPU dispatch is not timed as work."""
    import torch
    from comfy import model_management

    device = model_management.get_torch_device()
    if device.type == "mps":
        torch.mps.synchronize()
    elif device.type == "cuda":
        torch.cuda.synchronize(device)
    elif device.type == "xpu":
        torch.xpu.synchronize(device)


def _current_timing_label() -> str | None:
    from comfy_execution.utils import get_executing_context

    context = get_executing_context()
    if context is None:
        return None
    from comfy_execution.progress import get_progress_state

    progress = get_progress_state()
    if progress.prompt_id != context.prompt_id or not progress.dynprompt.has_node(context.node_id):
        return None
    metadata = progress.dynprompt.get_node(context.node_id).get("_meta", {})
    label = metadata.get("easy_media_timing") if isinstance(metadata, dict) else None
    return label if isinstance(label, str) and label else None


def instrument_node_timing(node_class: type) -> None:
    """Wrap a native method once; only runtime-tagged project nodes are timed.

    Preserve the descriptor and forward the actual cls/self so V3 execution
    clones keep their hidden inputs. No execution-engine hooks or node registry
    replacements are required; timing metadata lives only in the dynamic graph.
    """
    method_name = "execute" if callable(getattr(node_class, "execute", None)) else node_class.FUNCTION
    descriptor = inspect.getattr_static(node_class, method_name)
    original = descriptor.__func__ if isinstance(descriptor, (classmethod, staticmethod)) else descriptor
    if getattr(original, "__easy_media_timed__", False):
        return

    if inspect.iscoroutinefunction(original):
        @wraps(original)
        async def timed(*args: Any, **kwargs: Any) -> Any:
            label = _current_timing_label()
            if label is None:
                return await original(*args, **kwargs)
            with log_stage_time("MultiTrack Project", label, synchronize=synchronize_execution_device):
                return await original(*args, **kwargs)
    else:
        @wraps(original)
        def timed(*args: Any, **kwargs: Any) -> Any:
            label = _current_timing_label()
            if label is None:
                return original(*args, **kwargs)
            with log_stage_time("MultiTrack Project", label, synchronize=synchronize_execution_device):
                return original(*args, **kwargs)

    timed.__easy_media_timed__ = True
    if isinstance(descriptor, classmethod):
        replacement = classmethod(timed)
    elif isinstance(descriptor, staticmethod):
        replacement = staticmethod(timed)
    else:
        replacement = timed
    setattr(node_class, method_name, replacement)


# ---------------------------------------------------------------------------
# Console privacy scrub
#
# The pack presents itself as "Moxie Multimedia Suite" and must not leak the
# underlying model/vendor names through console output. The main offender is
# ComfyUI core itself (e.g. model_patcher logs "Model MiniMaxH3 prepared for
# dynamic VRAM loading ..." using the detected architecture class name).
#
# The filter below rewrites those tokens in DISPLAYED text only (record.msg,
# and pre-rendered tracebacks when they contain a token); no execution logic,
# node registration or workflow behaviour is affected.
# ---------------------------------------------------------------------------

# The cloud-API host must stay readable for provider debugging.
_SCRUB_PROTECT = ("minimaxi.com", "\x00MOXIEPROTECT\x00")

_SCRUB_PAIRS = (
    ("MiniMaxH3", "Moxie"),
    ("MiniMax H3", "Moxie"),
    ("MINIMAX_H3", "MOXIE"),
    ("minimax_h3", "moxie"),
    ("MiniMax", "Moxie"),
    ("MINIMAX", "MOXIE"),
    ("minimax", "moxie"),
)
_SCRUB_TOKENS = tuple(old for old, _ in _SCRUB_PAIRS)


def _scrub_text(text: str) -> str:
    protected, sentinel = _SCRUB_PROTECT
    text = text.replace(protected, sentinel)
    for old, new in _SCRUB_PAIRS:
        text = text.replace(old, new)
    return text.replace(sentinel, protected)


class _ConsolePrivacyFilter(logging.Filter):
    """Rewrite internal model/vendor tokens before records reach the console."""

    def filter(self, record: logging.LogRecord) -> bool:  # noqa: A003
        try:
            message = record.getMessage()
        except Exception:
            message = None
        if isinstance(message, str):
            scrubbed = _scrub_text(message)
            if scrubbed != message:
                record.msg = scrubbed
                record.args = None
        if record.exc_info is not None and record.exc_text is None:
            import traceback

            try:
                rendered = "".join(traceback.format_exception(*record.exc_info))
            except Exception:
                rendered = None
            if rendered is not None and any(token in rendered for token in _SCRUB_TOKENS):
                record.exc_text = _scrub_text(rendered)
                record.exc_info = None
        return True


_privacy_filter_installed = False


def install_console_privacy_filter() -> None:
    """Attach the console scrub to the root logger and its current handlers.

    Idempotent. ComfyUI installs its own root handlers during startup, before
    custom node packs are imported, so attaching here covers everything the
    user sees in the console afterwards."""
    global _privacy_filter_installed
    if _privacy_filter_installed:
        return
    _privacy_filter_installed = True
    root = logging.getLogger()
    scrubber = _ConsolePrivacyFilter()
    root.addFilter(scrubber)
    for handler in list(root.handlers):
        handler.addFilter(scrubber)