alexds9
77bcb3dc2e
Added output info for: ckpt_path, saved_metadata_str, saved_prompt_str, saved_extra_str.
233 lines
9.2 KiB
Python
233 lines
9.2 KiB
Python
import json
|
|
import os
|
|
import folder_paths
|
|
import comfy.sd
|
|
|
|
class SaveCheckpointWithMetadata:
|
|
"""
|
|
Save a checkpoint with explicit control over safetensors metadata.
|
|
|
|
Modes:
|
|
- replace: write ONLY metadata_json. Ignore prompt_override and EXTRA_PNGINFO.
|
|
- merge_minimal: base = prompt (override or hidden PROMPT) plus EXTRA_PNGINFO (if include_extra_pnginfo=True),
|
|
then overlay metadata_json (your keys win).
|
|
|
|
File naming modes:
|
|
- smart_counter: if unsuffixed file does not exist, use it; otherwise continue from next free index
|
|
based on existing files (never backwards).
|
|
- no_counter_overwrite: always use unsuffixed file name; overwrite if it exists.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL", {
|
|
"tooltip": "MODEL to serialize into a .safetensors checkpoint. Wire your merged/trained model here."
|
|
}),
|
|
"filename_prefix": ("STRING", {
|
|
"default": "checkpoints/CustomMeta",
|
|
"tooltip": "Subfolder/prefix under output directory. Example: checkpoints/MyModel."
|
|
}),
|
|
"filename_mode": (["smart_counter", "no_counter_overwrite"], {
|
|
"default": "smart_counter",
|
|
"tooltip": "smart_counter: first save uses prefix.safetensors if free, later saves use prefix_00001_.safetensors, ... no_counter_overwrite: always write to prefix.safetensors (overwrite if exists)."
|
|
}),
|
|
"metadata_json": ("STRING", {
|
|
"default": "{}",
|
|
"multiline": True,
|
|
"tooltip": "JSON object of header keys to write. Values must be strings; non-strings are JSON-encoded for you. Example: {\"author\":\"Alex\",\"training\":{\"epochs\":80}}"
|
|
}),
|
|
"metadata_mode": (["replace", "merge_minimal"], {
|
|
"default": "replace",
|
|
"tooltip": "replace: write only metadata_json. merge_minimal: start with prompt and optional EXTRA_PNGINFO, then apply metadata_json (your keys win)."
|
|
}),
|
|
"include_extra_pnginfo": ("BOOLEAN", {
|
|
"default": True,
|
|
"tooltip": "Merge mode only. When ON, copy keys from hidden EXTRA_PNGINFO into the header before applying metadata_json."
|
|
}),
|
|
"prompt_override": ("STRING", {
|
|
"default": "",
|
|
"multiline": True,
|
|
"tooltip": "Optional override for the hidden PROMPT. Ignored in replace. In merge_minimal, becomes the base prompt unless you also set prompt in metadata_json."
|
|
}),
|
|
},
|
|
"optional": {
|
|
"clip": ("CLIP", {
|
|
"tooltip": "Optional CLIP to embed into the saved checkpoint."
|
|
}),
|
|
"vae": ("VAE", {
|
|
"tooltip": "Optional VAE to embed into the saved checkpoint."
|
|
}),
|
|
"clip_vision": ("CLIP_VISION", {
|
|
"tooltip": "Optional CLIP_VISION to embed into the saved checkpoint."
|
|
}),
|
|
},
|
|
"hidden": {
|
|
"prompt": "PROMPT",
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
},
|
|
}
|
|
|
|
# Text outputs to inspect what was written
|
|
RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING")
|
|
RETURN_NAMES = ("ckpt_path", "saved_metadata", "saved_prompt", "saved_extra_pnginfo")
|
|
FUNCTION = "save"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "advanced/model_merging"
|
|
|
|
def _coerce_metadata(self, dct):
|
|
if dct is None:
|
|
return {}
|
|
if not isinstance(dct, dict):
|
|
raise ValueError("metadata_json must be a JSON object.")
|
|
out = {}
|
|
for k, v in dct.items():
|
|
ks = str(k)
|
|
out[ks] = v if isinstance(v, str) else json.dumps(v, ensure_ascii=False)
|
|
return out
|
|
|
|
def _serialize_prompt(self, prompt_obj_or_str):
|
|
if prompt_obj_or_str is None:
|
|
return None
|
|
if isinstance(prompt_obj_or_str, str):
|
|
s = prompt_obj_or_str.strip()
|
|
return s if s else None
|
|
return json.dumps(prompt_obj_or_str, ensure_ascii=False)
|
|
|
|
def _build_unsuffixed_path(self, full_dir, base):
|
|
return os.path.join(full_dir, f"{base}.safetensors")
|
|
|
|
def _build_suffixed_path(self, full_dir, base, n):
|
|
return os.path.join(full_dir, f"{base}_{n:05}_.safetensors")
|
|
|
|
def _next_suffix_from_dir(self, full_dir, base):
|
|
"""
|
|
Inspect existing files and return the next numeric suffix.
|
|
If no suffixed files exist, return 1.
|
|
Looks for: base_00001_.safetensors style names only.
|
|
"""
|
|
next_n = 1
|
|
try:
|
|
for name in os.listdir(full_dir):
|
|
if not name.startswith(base + "_") or not name.endswith("_.safetensors"):
|
|
continue
|
|
middle = name[len(base) + 1:-len("_.safetensors")]
|
|
if middle.isdigit():
|
|
n = int(middle)
|
|
if n >= next_n:
|
|
next_n = n + 1
|
|
except FileNotFoundError:
|
|
# Directory does not exist yet; caller ensures it will be created
|
|
next_n = 1
|
|
return next_n
|
|
|
|
def _choose_ckpt_path(self, full_dir, base, filename_mode):
|
|
"""
|
|
Filename policy:
|
|
- smart_counter:
|
|
If base.safetensors does not exist, use it.
|
|
Else compute the next suffix from existing files and use that.
|
|
- no_counter_overwrite:
|
|
Always use base.safetensors (overwrite if exists).
|
|
"""
|
|
unsuffixed = self._build_unsuffixed_path(full_dir, base)
|
|
|
|
if filename_mode == "no_counter_overwrite":
|
|
return unsuffixed
|
|
|
|
# smart_counter
|
|
if not os.path.exists(unsuffixed):
|
|
return unsuffixed
|
|
|
|
start_n = self._next_suffix_from_dir(full_dir, base)
|
|
path = self._build_suffixed_path(full_dir, base, start_n)
|
|
while os.path.exists(path):
|
|
start_n += 1
|
|
path = self._build_suffixed_path(full_dir, base, start_n)
|
|
return path
|
|
|
|
def save(self,
|
|
model,
|
|
filename_prefix,
|
|
filename_mode,
|
|
metadata_json,
|
|
metadata_mode,
|
|
include_extra_pnginfo,
|
|
prompt_override,
|
|
clip=None,
|
|
vae=None,
|
|
clip_vision=None,
|
|
prompt=None,
|
|
extra_pnginfo=None):
|
|
|
|
# Parse user metadata
|
|
try:
|
|
user_meta_obj = json.loads(metadata_json) if metadata_json.strip() else {}
|
|
except Exception as e:
|
|
raise ValueError(f"Invalid metadata_json: {e}")
|
|
user_meta = self._coerce_metadata(user_meta_obj)
|
|
|
|
# Effective prompt
|
|
if prompt_override and prompt_override.strip():
|
|
try:
|
|
parsed = json.loads(prompt_override)
|
|
effective_prompt = self._serialize_prompt(parsed)
|
|
except Exception:
|
|
effective_prompt = self._serialize_prompt(prompt_override)
|
|
else:
|
|
effective_prompt = self._serialize_prompt(prompt)
|
|
|
|
# Base metadata by mode, track included EXTRA_PNGINFO keys
|
|
base_metadata = {}
|
|
saved_extra_subset = {}
|
|
if metadata_mode == "merge_minimal":
|
|
if effective_prompt is not None:
|
|
base_metadata["prompt"] = effective_prompt
|
|
if include_extra_pnginfo and isinstance(extra_pnginfo, dict):
|
|
for k, v in extra_pnginfo.items():
|
|
try:
|
|
val = v if isinstance(v, str) else json.dumps(v, ensure_ascii=False)
|
|
except Exception:
|
|
val = str(v)
|
|
key = str(k)
|
|
base_metadata[key] = val
|
|
saved_extra_subset[key] = val
|
|
|
|
# Final metadata
|
|
final_meta = base_metadata.copy()
|
|
final_meta.update(user_meta)
|
|
|
|
# Resolve output path and ensure directory exists
|
|
out_dir = folder_paths.get_output_directory()
|
|
full_dir, base, _counter_hint, _subfolder, _ = folder_paths.get_save_image_path(filename_prefix, out_dir)
|
|
os.makedirs(full_dir, exist_ok=True)
|
|
|
|
# Choose final path per selected filename_mode
|
|
ckpt_path = self._choose_ckpt_path(full_dir, base, filename_mode)
|
|
|
|
# Save checkpoint
|
|
comfy.sd.save_checkpoint(
|
|
ckpt_path,
|
|
model,
|
|
clip=clip,
|
|
vae=vae,
|
|
clip_vision=clip_vision,
|
|
metadata=final_meta,
|
|
extra_keys={}
|
|
)
|
|
|
|
# Prepare text outputs
|
|
saved_metadata_str = json.dumps(final_meta, ensure_ascii=False, indent=2)
|
|
saved_prompt_str = final_meta.get("prompt", "") or ""
|
|
saved_extra_str = json.dumps(saved_extra_subset, ensure_ascii=False, indent=2)
|
|
|
|
return (ckpt_path, saved_metadata_str, saved_prompt_str, saved_extra_str)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SaveCheckpointWithMetadata": SaveCheckpointWithMetadata,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Save Checkpoint with Metadata": "Save Checkpoint with Metadata",
|
|
}
|