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.
This commit is contained in:
Will Miao
2025-11-14 15:46:11 +08:00
parent 3d5ec4a9f1
commit 7850131969

View File

@@ -72,6 +72,18 @@ class GGUFLoaderExtractor(NodeMetadataExtractor):
model_name = inputs.get("gguf_name") model_name = inputs.get("gguf_name")
_store_checkpoint_metadata(metadata, node_id, model_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): class TSCCheckpointLoaderExtractor(NodeMetadataExtractor):
@staticmethod @staticmethod
def extract(node_id, inputs, outputs, metadata): def extract(node_id, inputs, outputs, metadata):
@@ -695,6 +707,8 @@ NODE_EXTRACTORS = {
"NunchakuQwenImageDiTLoader": NunchakuQwenImageDiTLoaderExtractor, # ComfyUI-Nunchaku "NunchakuQwenImageDiTLoader": NunchakuQwenImageDiTLoaderExtractor, # ComfyUI-Nunchaku
"LoaderGGUF": GGUFLoaderExtractor, # calcuis gguf "LoaderGGUF": GGUFLoaderExtractor, # calcuis gguf
"LoaderGGUFAdvanced": GGUFLoaderExtractor, # calcuis gguf "LoaderGGUFAdvanced": GGUFLoaderExtractor, # calcuis gguf
"GGUFLoaderKJ": KJNodesModelLoaderExtractor, # KJNodes
"DiffusionModelLoaderKJ": KJNodesModelLoaderExtractor, # KJNodes
"UNETLoader": UNETLoaderExtractor, # Updated to use dedicated extractor "UNETLoader": UNETLoaderExtractor, # Updated to use dedicated extractor
"UnetLoaderGGUF": UNETLoaderExtractor, # Updated to use dedicated extractor "UnetLoaderGGUF": UNETLoaderExtractor, # Updated to use dedicated extractor
"LoraLoader": LoraLoaderExtractor, "LoraLoader": LoraLoaderExtractor,