From 8158441a925268d47a3f63c8f019f7dd1b52c2a8 Mon Sep 17 00:00:00 2001
From: Will Miao <willmiao.dev@gmail.com>
Date: Mon, 17 Nov 2025 08:52:51 +0800
Subject: [PATCH] feat: add CheckpointLoaderKJ extractor and improve model
 filename handling, fixes #666

- Add CheckpointLoaderKJ to NODE_EXTRACTORS mapping for KJNodes support
- Enhance model filename generation in SaveImage to handle different data types
- Add proper type checking and fallback for model metadata values
- Improve robustness when processing checkpoint paths for filename generation
---
 py/metadata_collector/node_extractors.py |  1 +
 py/nodes/save_image.py                   | 14 ++++++++++----
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/py/metadata_collector/node_extractors.py b/py/metadata_collector/node_extractors.py
index 3c7b868b..01768c10 100644
--- a/py/metadata_collector/node_extractors.py
+++ b/py/metadata_collector/node_extractors.py
@@ -709,6 +709,7 @@ NODE_EXTRACTORS = {
     "LoaderGGUFAdvanced": GGUFLoaderExtractor,  # calcuis gguf
     "GGUFLoaderKJ": KJNodesModelLoaderExtractor,  # KJNodes
     "DiffusionModelLoaderKJ": KJNodesModelLoaderExtractor,  # KJNodes
+    "CheckpointLoaderKJ": CheckpointLoaderExtractor,  # KJNodes
     "UNETLoader": UNETLoaderExtractor,          # Updated to use dedicated extractor
     "UnetLoaderGGUF": UNETLoaderExtractor,  # Updated to use dedicated extractor
     "LoraLoader": LoraLoaderExtractor,
diff --git a/py/nodes/save_image.py b/py/nodes/save_image.py
index cb3a8e9e..dbf44d07 100644
--- a/py/nodes/save_image.py
+++ b/py/nodes/save_image.py
@@ -273,9 +273,15 @@ class SaveImage:
                     length = int(parts[1])
                     prompt = prompt[:length]
                 filename = filename.replace(segment, prompt.strip())
-            elif key == "model" and 'checkpoint' in metadata_dict:
-                model = metadata_dict.get('checkpoint', '')
-                model = os.path.splitext(os.path.basename(model))[0]
+            elif key == "model":
+                model_value = metadata_dict.get('checkpoint')
+                if isinstance(model_value, (bytes, os.PathLike)):
+                    model_value = str(model_value)
+
+                if not isinstance(model_value, str) or not model_value:
+                    model = "model_unavailable"
+                else:
+                    model = os.path.splitext(os.path.basename(model_value))[0]
                 if len(parts) >= 2:
                     length = int(parts[1])
                     model = model[:length]
@@ -442,4 +448,4 @@ class SaveImage:
             add_counter_to_filename
         )
         
-        return (images,)
\ No newline at end of file
+        return (images,)