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
This commit is contained in:
Will Miao
2025-11-17 08:52:51 +08:00
parent 5600471093
commit 8158441a92
2 changed files with 11 additions and 4 deletions

View File

@@ -709,6 +709,7 @@ NODE_EXTRACTORS = {
"LoaderGGUFAdvanced": GGUFLoaderExtractor, # calcuis gguf "LoaderGGUFAdvanced": GGUFLoaderExtractor, # calcuis gguf
"GGUFLoaderKJ": KJNodesModelLoaderExtractor, # KJNodes "GGUFLoaderKJ": KJNodesModelLoaderExtractor, # KJNodes
"DiffusionModelLoaderKJ": KJNodesModelLoaderExtractor, # KJNodes "DiffusionModelLoaderKJ": KJNodesModelLoaderExtractor, # KJNodes
"CheckpointLoaderKJ": CheckpointLoaderExtractor, # 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,

View File

@@ -273,9 +273,15 @@ class SaveImage:
length = int(parts[1]) length = int(parts[1])
prompt = prompt[:length] prompt = prompt[:length]
filename = filename.replace(segment, prompt.strip()) filename = filename.replace(segment, prompt.strip())
elif key == "model" and 'checkpoint' in metadata_dict: elif key == "model":
model = metadata_dict.get('checkpoint', '') model_value = metadata_dict.get('checkpoint')
model = os.path.splitext(os.path.basename(model))[0] 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: if len(parts) >= 2:
length = int(parts[1]) length = int(parts[1])
model = model[:length] model = model[:length]