128 lines
5.9 KiB
Python
Executable File
128 lines
5.9 KiB
Python
Executable File
import json
|
|
import os
|
|
from safetensors import safe_open
|
|
|
|
# Node class mapping and display name (how it appears in ComfyUI)
|
|
NODE_CLASS_MAPPINGS = {}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
|
|
class EmAySee_LoRAMetadataExtractor: # Prepend EmAySee_ to the class name
|
|
"""
|
|
A ComfyUI custom node to extract and display training metadata
|
|
from a .safetensors LoRA file.
|
|
|
|
It attempts to find the full Kohya_ss configuration under '__metadata__',
|
|
or extracts individual 'ss_' prefixed keys if '__metadata__' is not present.
|
|
"""
|
|
def __init__(self):
|
|
pass
|
|
|
|
# Define the return types of the node's outputs
|
|
# This node will output a single STRING
|
|
RETURN_TYPES = ("STRING",)
|
|
|
|
# Define the names of the node's outputs
|
|
RETURN_NAMES = ("metadata_json_string",)
|
|
|
|
# Define the category where the node will appear in ComfyUI
|
|
FUNCTION = "extract_metadata"
|
|
CATEGORY = "Mac Custom Nodes/LoRA Tools" # You can change this category if you prefer
|
|
|
|
# Define the input types of the node
|
|
# This node requires a file path to a .safetensors file
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"lora_filepath": ("STRING", {"multiline": False, "default": ""}),
|
|
},
|
|
"optional": {
|
|
# Optional input for a file upload, if you prefer to drag-and-drop
|
|
# Note: ComfyUI's file input might not directly give a full path easily
|
|
# Manual string path is often more reliable for existing files.
|
|
"lora_upload": ("LORA", {"forceInput": True}), # This allows drag-and-drop, but you'd still need to get the path
|
|
}
|
|
}
|
|
|
|
def extract_metadata(self, lora_filepath, lora_upload=None):
|
|
"""
|
|
Extracts metadata from the specified .safetensors file.
|
|
|
|
Args:
|
|
lora_filepath (str): The full path to the .safetensors LoRA file.
|
|
lora_upload (str, optional): ComfyUI's internal representation of an uploaded LORA.
|
|
Not directly used for path here, but included for compatibility.
|
|
|
|
Returns:
|
|
tuple: A tuple containing the metadata as a formatted JSON string.
|
|
"""
|
|
if not lora_filepath:
|
|
return ("Error: LoRA filepath cannot be empty.",)
|
|
|
|
# Ensure the file exists
|
|
if not os.path.exists(lora_filepath):
|
|
return (f"Error: File not found at '{lora_filepath}'. Please check the path.",)
|
|
if not os.path.isfile(lora_filepath):
|
|
return (f"Error: Path '{lora_filepath}' is not a valid file.",)
|
|
|
|
try:
|
|
with safe_open(lora_filepath, framework="pt", device="cpu") as f:
|
|
metadata = f.metadata()
|
|
output_data = {}
|
|
|
|
# Attempt to get the full Kohya_ss config if nested under __metadata__ (newer Kohya)
|
|
if "__metadata__" in metadata:
|
|
try:
|
|
output_data = json.loads(metadata["__metadata__"])
|
|
# Add a flag to indicate it's the full config
|
|
output_data["_metadata_source"] = "Full Kohya_ss Config (__metadata__)"
|
|
except json.JSONDecodeError:
|
|
output_data["_error"] = "Failed to parse __metadata__ as JSON."
|
|
output_data["_raw_metadata"] = metadata["__metadata__"]
|
|
output_data["_metadata_source"] = "Raw __metadata__ (Parse Error)"
|
|
else:
|
|
# If not, extract individual ss_ prefixed keys and their values
|
|
extracted_params = {}
|
|
for key, value in metadata.items():
|
|
# First, attempt general type conversion (int, float, bool, None)
|
|
if isinstance(value, str):
|
|
if value.lower() == "true":
|
|
value = True
|
|
elif value.lower() == "false":
|
|
value = False
|
|
elif value.lower() == "none": # Handle "None" as Python None
|
|
value = None
|
|
else:
|
|
try:
|
|
if "." in value: # Likely a float
|
|
value = float(value)
|
|
else: # Try integer
|
|
value = int(value)
|
|
except ValueError:
|
|
pass # Keep as string if not a simple number/boolean
|
|
|
|
# Then, specifically parse known nested JSON strings
|
|
if key in ["ss_tag_frequency", "ss_bucket_info", "ss_datasets"]:
|
|
if isinstance(value, str): # Only try to parse if it's still a string
|
|
try:
|
|
value = json.loads(value)
|
|
except json.JSONDecodeError:
|
|
# If parsing fails, keep it as a string but note the error
|
|
value = f"ERROR_PARSING_JSON: {value}"
|
|
|
|
# Store the processed value
|
|
extracted_params[key] = value
|
|
|
|
output_data = extracted_params
|
|
output_data["_metadata_source"] = "Extracted ss_ Prefixed Keys"
|
|
|
|
# Return the data as a pretty-printed JSON string
|
|
return (json.dumps(output_data, indent=2),)
|
|
|
|
except Exception as e:
|
|
return (f"An unexpected error occurred during metadata extraction: {e}",)
|
|
|
|
# Register the node
|
|
NODE_CLASS_MAPPINGS["EmAySee_LoRAMetadataExtractor"] = EmAySee_LoRAMetadataExtractor # Prepend EmAySee_
|
|
NODE_DISPLAY_NAME_MAPPINGS["EmAySee_LoRA Metadata Extractor"] = "EmAySee_LoRA Metadata Extractor" # Prepend EmAySee_
|