369 lines
13 KiB
Python
369 lines
13 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.
|
|
|
|
File naming modes:
|
|
- smart_counter: if unsuffixed file does not exist, use it; otherwise continue
|
|
from next free index based on existing files.
|
|
- no_counter_overwrite: always use unsuffixed file name; overwrite if it exists.
|
|
"""
|
|
|
|
SEARCH_ALIASES = ["save checkpoint metadata", "checkpoint metadata", "safetensors metadata"]
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
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 and overwrite if it 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."
|
|
}),
|
|
"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 mode. In merge_minimal mode, becomes the base prompt unless metadata_json also contains prompt."
|
|
}),
|
|
},
|
|
"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",
|
|
},
|
|
}
|
|
|
|
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():
|
|
key = str(k)
|
|
out[key] = 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
|
|
"""
|
|
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:
|
|
next_n = 1
|
|
|
|
return next_n
|
|
|
|
def _choose_save_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.
|
|
"""
|
|
unsuffixed = self._build_unsuffixed_path(full_dir, base)
|
|
|
|
if filename_mode == "no_counter_overwrite":
|
|
return unsuffixed
|
|
|
|
if not os.path.exists(unsuffixed):
|
|
return unsuffixed
|
|
|
|
suffix_n = self._next_suffix_from_dir(full_dir, base)
|
|
save_path = self._build_suffixed_path(full_dir, base, suffix_n)
|
|
|
|
while os.path.exists(save_path):
|
|
suffix_n += 1
|
|
save_path = self._build_suffixed_path(full_dir, base, suffix_n)
|
|
|
|
return save_path
|
|
|
|
def _save_file(self, save_path, model, clip, vae, clip_vision, final_meta):
|
|
"""
|
|
Base checkpoint save implementation.
|
|
|
|
This is used by SaveCheckpointWithMetadata and can be overridden by
|
|
subclasses that need different save behavior.
|
|
"""
|
|
comfy.sd.save_checkpoint(
|
|
save_path,
|
|
model,
|
|
clip=clip,
|
|
vae=vae,
|
|
clip_vision=clip_vision,
|
|
metadata=final_meta,
|
|
extra_keys={}
|
|
)
|
|
|
|
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,
|
|
):
|
|
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)
|
|
|
|
if prompt_override and prompt_override.strip():
|
|
try:
|
|
parsed_prompt_override = json.loads(prompt_override)
|
|
effective_prompt = self._serialize_prompt(parsed_prompt_override)
|
|
except Exception:
|
|
effective_prompt = self._serialize_prompt(prompt_override)
|
|
else:
|
|
effective_prompt = self._serialize_prompt(prompt)
|
|
|
|
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():
|
|
key = str(k)
|
|
|
|
try:
|
|
val = v if isinstance(v, str) else json.dumps(v, ensure_ascii=False)
|
|
except Exception:
|
|
val = str(v)
|
|
|
|
base_metadata[key] = val
|
|
saved_extra_subset[key] = val
|
|
|
|
final_meta = base_metadata.copy()
|
|
final_meta.update(user_meta)
|
|
|
|
out_dir = folder_paths.get_output_directory()
|
|
full_dir, base, _counter_hint, _subfolder, _filename = folder_paths.get_save_image_path(
|
|
filename_prefix,
|
|
out_dir
|
|
)
|
|
os.makedirs(full_dir, exist_ok=True)
|
|
|
|
save_path = self._choose_save_path(full_dir, base, filename_mode)
|
|
|
|
self._save_file(
|
|
save_path=save_path,
|
|
model=model,
|
|
clip=clip,
|
|
vae=vae,
|
|
clip_vision=clip_vision,
|
|
final_meta=final_meta,
|
|
)
|
|
|
|
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 (save_path, saved_metadata_str, saved_prompt_str, saved_extra_str)
|
|
|
|
|
|
class SaveDiffusionModelWithMetadata(SaveCheckpointWithMetadata):
|
|
"""
|
|
Save only the MODEL/diffusion model with explicit safetensors metadata control.
|
|
|
|
This shares metadata, filename replacement, and counter behavior with
|
|
SaveCheckpointWithMetadata, but exposes no CLIP/VAE inputs.
|
|
"""
|
|
|
|
SEARCH_ALIASES = ["export model", "save diffusion model", "model save metadata"]
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL", {
|
|
"tooltip": "MODEL to serialize into a .safetensors diffusion model file."
|
|
}),
|
|
"filename_prefix": ("STRING", {
|
|
"default": "diffusion_models/CustomMeta",
|
|
"tooltip": "Subfolder/prefix under output directory. Example: diffusion_models/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: overwrite prefix.safetensors."
|
|
}),
|
|
"metadata_json": ("STRING", {
|
|
"default": "{}",
|
|
"multiline": True,
|
|
"tooltip": "JSON object of safetensors header keys to write."
|
|
}),
|
|
"metadata_mode": (["replace", "merge_minimal"], {
|
|
"default": "replace",
|
|
"tooltip": "replace: write only metadata_json. merge_minimal: include prompt/extra_pnginfo, then apply metadata_json."
|
|
}),
|
|
"include_extra_pnginfo": ("BOOLEAN", {
|
|
"default": True,
|
|
"tooltip": "Merge mode only. Copy hidden EXTRA_PNGINFO keys before applying metadata_json."
|
|
}),
|
|
"prompt_override": ("STRING", {
|
|
"default": "",
|
|
"multiline": True,
|
|
"tooltip": "Optional replacement for hidden PROMPT in merge_minimal mode."
|
|
}),
|
|
},
|
|
"hidden": {
|
|
"prompt": "PROMPT",
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING")
|
|
RETURN_NAMES = ("model_path", "saved_metadata", "saved_prompt", "saved_extra_pnginfo")
|
|
FUNCTION = "save"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "advanced/model_merging"
|
|
|
|
def _save_file(self, save_path, model, clip, vae, clip_vision, final_meta):
|
|
"""
|
|
Diffusion-model-only save.
|
|
|
|
This intentionally ignores CLIP/VAE/CLIP_VISION inputs because this node
|
|
only exposes and saves the MODEL object.
|
|
"""
|
|
comfy.sd.save_checkpoint(
|
|
save_path,
|
|
model,
|
|
clip=None,
|
|
vae=None,
|
|
clip_vision=None,
|
|
metadata=final_meta,
|
|
extra_keys={}
|
|
)
|
|
|
|
def save(
|
|
self,
|
|
model,
|
|
filename_prefix,
|
|
filename_mode,
|
|
metadata_json,
|
|
metadata_mode,
|
|
include_extra_pnginfo,
|
|
prompt_override,
|
|
prompt=None,
|
|
extra_pnginfo=None,
|
|
):
|
|
return super().save(
|
|
model=model,
|
|
filename_prefix=filename_prefix,
|
|
filename_mode=filename_mode,
|
|
metadata_json=metadata_json,
|
|
metadata_mode=metadata_mode,
|
|
include_extra_pnginfo=include_extra_pnginfo,
|
|
prompt_override=prompt_override,
|
|
clip=None,
|
|
vae=None,
|
|
clip_vision=None,
|
|
prompt=prompt,
|
|
extra_pnginfo=extra_pnginfo,
|
|
)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SaveCheckpointWithMetadata": SaveCheckpointWithMetadata,
|
|
"SaveDiffusionModelWithMetadata": SaveDiffusionModelWithMetadata,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SaveCheckpointWithMetadata": "Save Checkpoint with Metadata",
|
|
"SaveDiffusionModelWithMetadata": "Save Diffusion Model with Metadata",
|
|
}
|