From 78501319698dbd1a26761a75e80d7e907e0cbb5a Mon Sep 17 00:00:00 2001
From: Will Miao <13051207myq@gmail.com>
Date: Fri, 14 Nov 2025 15:46:11 +0800
Subject: [PATCH] feat: add metadata extractor for KJNodes model loaders, see
 #666

Add KJNodesModelLoaderExtractor to handle metadata extraction from KJNodes loaders that expose model_name. This supports GGUFLoaderKJ and DiffusionModelLoaderKJ nodes, ensuring consistent checkpoint metadata collection across different node types.
---
 py/metadata_collector/node_extractors.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/py/metadata_collector/node_extractors.py b/py/metadata_collector/node_extractors.py
index e759447e..3c7b868b 100644
--- a/py/metadata_collector/node_extractors.py
+++ b/py/metadata_collector/node_extractors.py
@@ -72,6 +72,18 @@ class GGUFLoaderExtractor(NodeMetadataExtractor):
         model_name = inputs.get("gguf_name")
         _store_checkpoint_metadata(metadata, node_id, model_name)
 
+
+class KJNodesModelLoaderExtractor(NodeMetadataExtractor):
+    """Extract metadata from KJNodes loaders that expose `model_name`."""
+
+    @staticmethod
+    def extract(node_id, inputs, outputs, metadata):
+        if not inputs or "model_name" not in inputs:
+            return
+
+        model_name = inputs.get("model_name")
+        _store_checkpoint_metadata(metadata, node_id, model_name)
+
 class TSCCheckpointLoaderExtractor(NodeMetadataExtractor):
     @staticmethod
     def extract(node_id, inputs, outputs, metadata):
@@ -695,6 +707,8 @@ NODE_EXTRACTORS = {
     "NunchakuQwenImageDiTLoader": NunchakuQwenImageDiTLoaderExtractor,  # ComfyUI-Nunchaku
     "LoaderGGUF": GGUFLoaderExtractor,  # calcuis gguf
     "LoaderGGUFAdvanced": GGUFLoaderExtractor,  # calcuis gguf
+    "GGUFLoaderKJ": KJNodesModelLoaderExtractor,  # KJNodes
+    "DiffusionModelLoaderKJ": KJNodesModelLoaderExtractor,  # KJNodes
     "UNETLoader": UNETLoaderExtractor,          # Updated to use dedicated extractor
     "UnetLoaderGGUF": UNETLoaderExtractor,  # Updated to use dedicated extractor
     "LoraLoader": LoraLoaderExtractor,