From ed39c5918a7f097f33d10e3925bb40186fff9527 Mon Sep 17 00:00:00 2001
From: Ethanfel <ethan.fel@ts-pc.fr>
Date: Thu, 9 Apr 2026 18:04:16 +0200
Subject: [PATCH] feat: extend prompt handler and add /nodes-stats/models
 endpoint

---
 __init__.py | 30 ++++++++++++++++++++++++++----
 1 file changed, 26 insertions(+), 4 deletions(-)

diff --git a/__init__.py b/__init__.py
index 4847872..b9d3c41 100644
--- a/__init__.py
+++ b/__init__.py
@@ -4,7 +4,7 @@ import threading
 from aiohttp import web
 from server import PromptServer
 
-from .mapper import NodePackageMapper
+from .mapper import NodePackageMapper, ModelMapper
 from .tracker import UsageTracker
 
 logger = logging.getLogger(__name__)
@@ -14,10 +14,11 @@ 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 records usage."""
+    """Called on every prompt submission. Extracts class_types and queues recording."""
     try:
         prompt = json_data.get("prompt", {})
         class_types = set()
@@ -26,9 +27,11 @@ def on_prompt_handler(json_data):
             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=tracker.record_usage,
-                args=(class_types, mapper),
+                target=_record_prompt,
+                args=(class_types, prompt),
                 daemon=True,
             ).start()
     except Exception:
@@ -36,6 +39,13 @@ def on_prompt_handler(json_data):
     return json_data
 
 
+def _record_prompt(class_types, prompt):
+    tracker.record_usage(class_types, mapper)
+    models = model_mapper.extract_models_from_prompt(prompt)
+    if models:
+        tracker.record_model_usage(models)
+
+
 PromptServer.instance.add_on_prompt_handler(on_prompt_handler)
 
 
@@ -62,11 +72,23 @@ async def get_node_stats(request):
         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)