nodes/a-l-e-x-d-s-9-ComfyUI-SaveCheckpointWithMetadata
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

v0.3.0

Browse Source 
Added output info for: ckpt_path, saved_metadata_str, saved_prompt_str, saved_extra_str.
This commit is contained in:
alexds9
2025-08-25 20:48:06 +03:00
parent 5d23a883c0
commit 77bcb3dc2e
3 changed files with 155 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+99 -4
View File
@@ -1,6 +1,6 @@
# ComfyUI-SaveCheckpointWithMetadata
Save .safetensors checkpoints with custom metadata and explicit filename behavior in ComfyUI.
Save `.safetensors` checkpoints with custom metadata and explicit filename behavior in ComfyUI.
- Write exactly the header you want via `metadata_json`.
- Optional `prompt_override` to replace the hidden PROMPT.
@@ -12,16 +12,111 @@ Save .safetensors checkpoints with custom metadata and explicit filename behavio
## Install
### ComfyUI Manager
- Manager -> Install via Git URL -> paste:
1) Manager -> Install via Git URL -> paste:
```
https://github.com/a-l-e-x-d-s-9/ComfyUI-SaveCheckpointWithMetadata.git
```
2) Restart ComfyUI.
- Restart ComfyUI.
If Manager shows a security level error, either install manually or add your repo to the Manager catalog. See Troubleshooting below.
### Manual
```bash
cd /path/to/ComfyUI/custom_nodes
git clone https://github.com/<your-account>/ComfyUI-SaveCheckpointWithMetadata.git
git clone https://github.com/a-l-e-x-d-s-9/ComfyUI-SaveCheckpointWithMetadata.git
# optional deps
python3 -m pip install -r ComfyUI-SaveCheckpointWithMetadata/requirements.txt
# optional post-install hook
python3 ComfyUI-SaveCheckpointWithMetadata/install.py || true
```
Restart ComfyUI.
## Node location
- Category: `advanced/model_merging`
- Display name: `Save Checkpoint with Metadata`
## Quick start
1) Wire your `MODEL` (and optionally `CLIP`, `VAE`, `CLIP_VISION`) into the node.
2) Choose `filename_prefix`, for example `checkpoints/MyModel`.
3) Pick `filename_mode`:
- `smart_counter` for automatic unique names.
- `no_counter_overwrite` to always write `prefix.safetensors`.
4) Pick `metadata_mode`:
- `replace` to write only your `metadata_json`.
- `merge_minimal` to start with prompt plus optional `EXTRA_PNGINFO`, then apply `metadata_json`.
5) Fill `metadata_json` with a JSON object of header keys.
6) Optionally set `prompt_override` to replace the hidden PROMPT in merge mode.
7) Queue. Your checkpoint is saved in ComfyUI output directory under the resolved subfolder.
## Inputs
| Input | Type | Description |
|------|------|-------------|
| model | MODEL | The model to serialize into a `.safetensors` checkpoint. |
| filename_prefix | STRING | Subfolder and base name under ComfyUI output directory, for example `checkpoints/MyModel`. |
| filename_mode | DROPDOWN | `smart_counter` or `no_counter_overwrite`. See Filename behavior. |
| metadata_json | STRING (multiline) | 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 | DROPDOWN | `replace`: write only `metadata_json`. `merge_minimal`: base header includes prompt and, optionally, `EXTRA_PNGINFO`, then your `metadata_json` overwrites or adds keys. |
| include_extra_pnginfo | BOOLEAN | Used only in `merge_minimal`. When on, copy keys from hidden `EXTRA_PNGINFO` into the base header. |
| prompt_override | STRING (multiline) | Optional override for the hidden PROMPT. Ignored in `replace`. In `merge_minimal`, becomes the base `prompt` unless you also set `prompt` in `metadata_json`. |
| clip | CLIP (optional) | Embed CLIP in the checkpoint. |
| vae | VAE (optional) | Embed VAE in the checkpoint. |
| clip_vision | CLIP_VISION (optional) | Embed CLIP_VISION in the checkpoint. |
Tip: To see tooltips, enable ComfyUI Settings -> Nodes -> Enable Tooltips.
## Filename behavior
- `smart_counter`
- If `prefix.safetensors` does not exist, write there.
- Otherwise continue with `prefix_00001_.safetensors`, `prefix_00002_.safetensors`, etc., starting from the next free index. If ComfyUI hints a large counter, we do not go backwards; we probe from `hint+1` upward.
- `no_counter_overwrite`
- Always write to `prefix.safetensors`. If it exists, it is overwritten.
## Examples
### Replace-only header
- metadata_mode: `replace`
- metadata_json:
```json
{"author":"AIxFun eStudio","project":"Flux-Kontext","modelspec.architecture":"sdxl"}
```
- Result: header contains exactly those keys.
### Merge prompt and extra info, then override
- metadata_mode: `merge_minimal`
- prompt_override: paste your workflow JSON or plain text
- include_extra_pnginfo: on
- metadata_json:
```json
{"prompt":"OVERRIDE","author":"Alex","training":{"epochs":80}}
```
## Notes
- Safetensors metadata requires `dict[str, str]`. Non-strings are JSON-encoded by the node.
- This node does not follow the stock `--disable-metadata` flag. It writes exactly what you specify.
- Hidden inputs `PROMPT` and `EXTRA_PNGINFO` are provided by ComfyUI runtime. The node only reads them when `merge_minimal` is selected.
## Troubleshooting
### Manager shows: This action is not allowed with this security level configuration.
- Option A: Manual install by `git clone` into `ComfyUI/custom_nodes/`.
- Option B: Lower Manager `security_level` to `middle` or `weak` in its `config.ini`, then restart ComfyUI.
- Option C: Submit your repo to the ComfyUI-Manager catalog (`custom-node-list.json`) so it can be installed at normal security. Until the PR is merged, users can use Install via Git URL or manual install.
### I see `prefix_00001_.safetensors` even on the first save
- In `smart_counter` we prefer `prefix.safetensors` if it does not exist. If you still get a numbered file, ensure the output directory exists and that no file named `prefix.safetensors` is present.
### I want modelspec.* keys like the stock node adds
- Add them to `metadata_json` explicitly. Example:
```json
{"modelspec.architecture":"sdxl","modelspec.vision":"clip"}
```
## License
No.
## Credits
Author: alexds9
Inspired by ComfyUI stock saving logic but with explicit metadata and naming control.
__init__.py
+55 -26
View File
@@ -13,9 +13,9 @@ class SaveCheckpointWithMetadata:
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,
never counting backwards even if counter_hint is large.
- no_counter_overwrite: always use unsuffixed file name; if it exists, overwrite it.
- 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
@@ -27,18 +27,17 @@ class SaveCheckpointWithMetadata:
}),
"filename_prefix": ("STRING", {
"default": "checkpoints/CustomMeta",
"tooltip": "Subfolder/prefix under output directory. Example: 'checkpoints/MyModel'."
"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', '..._00002_...', etc. no_counter_overwrite: always save to 'prefix.safetensors' and overwrite if it exists."
"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}}"
}),
# Dropdown: list-of-choices renders as a combo box
"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)."
@@ -50,7 +49,7 @@ class SaveCheckpointWithMetadata:
"prompt_override": ("STRING", {
"default": "",
"multiline": True,
"tooltip": "Optional override for the hidden PROMPT. In replace mode this is ignored. In merge_minimal, becomes the base 'prompt' unless you also set 'prompt' in metadata_json."
"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": {
@@ -65,13 +64,14 @@ class SaveCheckpointWithMetadata:
}),
},
"hidden": {
# Hidden inputs provided by ComfyUI runtime
"prompt": "PROMPT",
"extra_pnginfo": "EXTRA_PNGINFO",
},
}
RETURN_TYPES = ()
# 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"
@@ -101,14 +101,35 @@ class SaveCheckpointWithMetadata:
def _build_suffixed_path(self, full_dir, base, n):
return os.path.join(full_dir, f"{base}_{n:05}_.safetensors")
def _choose_ckpt_path(self, full_dir, base, counter_hint, filename_mode):
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, continue from next free index starting at max(counter_hint, 1) + 1, scanning upward.
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 it exists).
Always use base.safetensors (overwrite if exists).
"""
unsuffixed = self._build_unsuffixed_path(full_dir, base)
@@ -119,11 +140,11 @@ class SaveCheckpointWithMetadata:
if not os.path.exists(unsuffixed):
return unsuffixed
n = max(counter_hint, 1) + 1
path = self._build_suffixed_path(full_dir, base, n)
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):
n += 1
path = self._build_suffixed_path(full_dir, base, n)
start_n += 1
path = self._build_suffixed_path(full_dir, base, start_n)
return path
def save(self,
@@ -157,31 +178,33 @@ class SaveCheckpointWithMetadata:
else:
effective_prompt = self._serialize_prompt(prompt)
# Base metadata by mode
# 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:
base_metadata[str(k)] = v if isinstance(v, str) else json.dumps(v, ensure_ascii=False)
val = v if isinstance(v, str) else json.dumps(v, ensure_ascii=False)
except Exception:
base_metadata[str(k)] = str(v)
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
# 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)
# Ensure directory exists before we choose the name
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, counter_hint, filename_mode)
ckpt_path = self._choose_ckpt_path(full_dir, base, filename_mode)
# Save checkpoint
comfy.sd.save_checkpoint(
@@ -193,7 +216,13 @@ class SaveCheckpointWithMetadata:
metadata=final_meta,
extra_keys={}
)
return {}
# 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,
pyproject.toml
+1 -1
View File
@@ -1,6 +1,6 @@
[project]
name = "save-checkpoint-with-metadata"
version = "0.2.0"
version = "0.3.0"
description = "ComfyUI custom node: save .safetensors with custom metadata and smart filename modes."
readme = "README.md"
requires-python = ">=3.10"