diff --git a/README.md b/README.md index 7e55e64..6ba7edc 100644 --- a/README.md +++ b/README.md @@ -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//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. diff --git a/__init__.py b/__init__.py index d785bc4..97b48b0 100644 --- a/__init__.py +++ b/__init__.py @@ -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, diff --git a/pyproject.toml b/pyproject.toml index b97701c..99bfc9a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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"