Comfyui-Nodes-Stats/__init__.py

import asyncio
import logging
import threading

from aiohttp import web
from server import PromptServer

from .mapper import NodePackageMapper, ModelMapper
from .node_introspect import find_disabled_pack_path, get_node_schema
from .tracker import UsageTracker

logger = logging.getLogger(__name__)

NODE_CLASS_MAPPINGS = {}
WEB_DIRECTORY = "./js"

mapper = NodePackageMapper()
tracker = UsageTracker()
model_mapper = ModelMapper()


def on_prompt_handler(json_data):
    """Called on every prompt submission. Extracts class_types and queues recording."""
    try:
        prompt = json_data.get("prompt", {})
        class_types = set()
        for node_id, node_data in prompt.items():
            ct = node_data.get("class_type")
            if ct:
                class_types.add(ct)
        if class_types:
            # Pass the full prompt to the thread — model extraction (which calls
            # INPUT_TYPES() on every node) happens off the main request thread.
            threading.Thread(
                target=_record_prompt,
                args=(class_types, prompt),
                daemon=True,
            ).start()
    except Exception:
        logger.warning("nodes-stats: error recording usage", exc_info=True)
    return json_data


def _record_prompt(class_types, prompt):
    try:
        tracker.record_usage(class_types, mapper)
    except Exception:
        logger.warning("nodes-stats: error recording node usage", exc_info=True)
    try:
        packages = {mapper.get_package(ct) for ct in class_types}
        packages.discard("__builtin__")
        packages.discard("__unknown__")
        tracker.reset_trials_for(packages)
    except Exception:
        logger.warning("nodes-stats: error resetting trials", exc_info=True)
    try:
        models = model_mapper.extract_models_from_prompt(prompt)
        if models:
            tracker.record_model_usage(models)
    except Exception:
        logger.warning("nodes-stats: error recording model usage", exc_info=True)


# Age temporary-enable trials once per process start (one "boot").
try:
    tracker.tick_boot_days()
except Exception:
    logger.warning("nodes-stats: error ticking trial boot days", exc_info=True)


PromptServer.instance.add_on_prompt_handler(on_prompt_handler)


routes = PromptServer.instance.routes


@routes.get("/nodes-stats/packages")
async def get_package_stats(request):
    try:
        stats = tracker.get_package_stats(mapper)
        return web.json_response(stats)
    except Exception:
        logger.error("nodes-stats: error getting package stats", exc_info=True)
        return web.json_response({"error": "internal error"}, status=500)


@routes.get("/nodes-stats/usage")
async def get_node_stats(request):
    try:
        stats = tracker.get_node_stats()
        return web.json_response(stats)
    except Exception:
        logger.error("nodes-stats: error getting node stats", exc_info=True)
        return web.json_response({"error": "internal error"}, status=500)


@routes.get("/nodes-stats/models")
async def get_model_stats(request):
    try:
        installed_by_type = model_mapper.get_all_models()
        stats = tracker.get_model_stats(installed_by_type)
        return web.json_response(stats)
    except Exception:
        logger.error("nodes-stats: error getting model stats", exc_info=True)
        return web.json_response({"error": "internal error"}, status=500)


@routes.post("/nodes-stats/reset")
async def reset_stats(request):
    try:
        tracker.reset()
        mapper.invalidate()
        model_mapper.invalidate()
        return web.json_response({"status": "ok"})
    except Exception:
        logger.error("nodes-stats: error resetting stats", exc_info=True)
        return web.json_response({"error": "internal error"}, status=500)


@routes.get("/nodes-stats/node-schema")
async def get_node_schema_route(request):
    """Parse a disabled pack's source and return one node's input/output schema.

    Read-only and never imports/executes the pack. Used by the mirror-search
    palette to draw a faithful node box for a node that isn't loaded.
    """
    try:
        class_type = request.query.get("class_type")
        pack = request.query.get("pack")
        if not class_type or not pack:
            return web.json_response({"error": "class_type and pack required"}, status=400)

        def _resolve():
            path = find_disabled_pack_path(pack)
            if not path:
                return {"parseable": False, "reason": "source_not_found"}
            return get_node_schema(class_type, path)

        # AST-parsing a whole pack can take tens of ms; keep it off the event loop.
        schema = await asyncio.get_event_loop().run_in_executor(None, _resolve)
        return web.json_response(schema)
    except Exception:
        logger.error("nodes-stats: error parsing node schema", exc_info=True)
        return web.json_response({"error": "internal error"}, status=500)


@routes.get("/nodes-stats/trials")
async def get_trials(request):
    try:
        return web.json_response(tracker.get_trials())
    except Exception:
        logger.error("nodes-stats: error getting trials", exc_info=True)
        return web.json_response({"error": "internal error"}, status=500)


@routes.post("/nodes-stats/trials/start")
async def start_trial(request):
    try:
        data = await request.json()
        package = data.get("package")
        if not package:
            return web.json_response({"error": "package required"}, status=400)
        tracker.start_trial(package)
        return web.json_response({"status": "ok"})
    except Exception:
        logger.error("nodes-stats: error starting trial", exc_info=True)
        return web.json_response({"error": "internal error"}, status=500)


@routes.post("/nodes-stats/trials/stop")
async def stop_trial(request):
    try:
        data = await request.json()
        package = data.get("package")
        if not package:
            return web.json_response({"error": "package required"}, status=400)
        tracker.stop_trial(package)
        return web.json_response({"status": "ok"})
    except Exception:
        logger.error("nodes-stats: error stopping trial", exc_info=True)
        return web.json_response({"error": "internal error"}, status=500)