diff --git a/nodes/civitai_metadata/filenames.py b/nodes/civitai_metadata/filenames.py new file mode 100644 index 0000000..050dec1 --- /dev/null +++ b/nodes/civitai_metadata/filenames.py @@ -0,0 +1,32 @@ +# Filename token expansion — part of ComfyUI-AI2Go-Utils. GPL-3.0. +# +# ComfyUI core's folder_paths.get_save_image_path expands %year%/%month%/%day%/... and subfolders, +# but NOT the %date:FORMAT% token the stock Save Image tooltip advertises (that lives in the web +# frontend and doesn't reach a custom node's widget). This expands %date:FORMAT% (and bare %date%) +# so filename prefixes like "%date:yyyy-MM-dd%/Krea-Turbo" work. Pure/comfy-free — the caller passes +# `now`; other tokens are left for get_save_image_path. +import re + +# .NET/luxon-style format tokens -> strftime. Longest-first so yyyy wins over yy in a single pass, +# and case distinguishes MM (month) from mm (minute), HH/hh (hour). +_FMT_TOKENS = [("yyyy", "%Y"), ("yy", "%y"), ("MM", "%m"), ("dd", "%d"), + ("HH", "%H"), ("hh", "%I"), ("mm", "%M"), ("ss", "%S")] +_FMT_MAP = dict(_FMT_TOKENS) +_FMT_RE = re.compile("|".join(re.escape(tok) for tok, _ in _FMT_TOKENS)) +_DATE_RE = re.compile(r"%date:([^%]*)%") + + +def _to_strftime(fmt): + return _FMT_RE.sub(lambda m: _FMT_MAP[m.group(0)], fmt) + + +def expand_date_tokens(prefix, now): + """Expand ``%date:FORMAT%`` (and bare ``%date%`` -> yyyy-MM-dd) in ``prefix`` using ``now``. + + ``now`` is a ``datetime``. Non-date tokens (``%year%``, ``%width%``, subfolders) are left + untouched for ``folder_paths.get_save_image_path``. + """ + if not prefix or "%date" not in prefix: + return prefix + prefix = _DATE_RE.sub(lambda m: now.strftime(_to_strftime(m.group(1))), prefix) + return prefix.replace("%date%", now.strftime("%Y-%m-%d")) diff --git a/nodes/civitai_metadata/tracer.py b/nodes/civitai_metadata/tracer.py index 165cd02..00448ec 100644 --- a/nodes/civitai_metadata/tracer.py +++ b/nodes/civitai_metadata/tracer.py @@ -14,7 +14,17 @@ from ..prompt_batch_core import parse_prompts, select_prompt SAMPLER_CLASSES = {"KSampler", "KSamplerAdvanced", "SamplerCustom", "SamplerCustomAdvanced"} CLIP_ENCODE_CLASSES = {"CLIPTextEncode"} BATCH_CLASS = "AI2GoPromptBatch" -CHECKPOINT_CLASSES = {"CheckpointLoaderSimple", "CheckpointLoader", "unCLIPCheckpointLoader"} +# Model-source loaders → (widget holding the filename, folder_paths folder to hash it from). +# Covers checkpoints AND standalone diffusion models (UNETLoader "Load Diffusion Model" — Flux/SD3/ +# Krea/etc.), plus common GGUF unet loaders. Unknown loaders fall through as pass-throughs. +MODEL_SOURCES = { + "CheckpointLoaderSimple": ("ckpt_name", "checkpoints"), + "CheckpointLoader": ("ckpt_name", "checkpoints"), + "unCLIPCheckpointLoader": ("ckpt_name", "checkpoints"), + "UNETLoader": ("unet_name", "diffusion_models"), + "UnetLoaderGGUF": ("unet_name", "diffusion_models"), + "UnetLoaderGGUFAdvanced": ("unet_name", "diffusion_models"), +} LORA_CLASSES = {"LoraLoader", "LoraLoaderModelOnly"} @@ -31,6 +41,7 @@ class TraceResult: clip_skip: object = None model_name: str = None model_file: str = None + model_folder: str = None loras: list = field(default_factory=list) unresolved: list = field(default_factory=list) @@ -190,11 +201,13 @@ def _trace_model_chain(prompt, sampler, r): nxt = ins.get("model") nid = str(nxt[0]) if _is_link(nxt) else None continue - if cls in CHECKPOINT_CLASSES: - ck = ins.get("ckpt_name") - if isinstance(ck, str): - r.model_name = _stem(ck) - r.model_file = ck + if cls in MODEL_SOURCES: + widget, folder = MODEL_SOURCES[cls] + name = ins.get(widget) + if isinstance(name, str): + r.model_name = _stem(name) + r.model_file = name + r.model_folder = folder break nxt = ins.get("model") # pass-through node (e.g. ModelSamplingDiscrete) nid = str(nxt[0]) if _is_link(nxt) else None diff --git a/nodes/save_civitai_metadata.py b/nodes/save_civitai_metadata.py index f87c464..9711147 100644 --- a/nodes/save_civitai_metadata.py +++ b/nodes/save_civitai_metadata.py @@ -6,6 +6,7 @@ # web/js/save_civitai_metadata.js (the Test button) — keep the two in sync. import os import logging +from datetime import datetime import numpy as np from PIL import Image @@ -24,6 +25,7 @@ from .civitai_metadata.sampler_names import to_a1111_sampler from .civitai_metadata.a1111 import format_parameters from .civitai_metadata.hashing import HashCache from .civitai_metadata.png_info import build_pnginfo +from .civitai_metadata.filenames import expand_date_tokens logger = logging.getLogger(__name__) @@ -49,7 +51,9 @@ def _save_with_metadata(cls, images, filename_prefix, save_workflow, overrides): # Resolve model + LoRA file paths to AutoV2 hashes (cached). cache = HashCache(_CACHE_PATH) - model_hash = cache.get(folder_paths.get_full_path("checkpoints", result.model_file)) if result.model_file else None + model_hash = cache.get( + folder_paths.get_full_path(result.model_folder or "checkpoints", result.model_file) + ) if result.model_file else None for lo in result.loras: path = folder_paths.get_full_path("loras", lo["file"]) if lo.get("file") else None lo["hash"] = cache.get(path) if path else None @@ -63,6 +67,9 @@ def _save_with_metadata(cls, images, filename_prefix, save_workflow, overrides): denoise=result.denoise, clip_skip=result.clip_skip, ) + # Expand %date:FORMAT% (frontend feature that doesn't reach a custom node's widget) before the + # rest of the tokens (%year%/subfolders) that get_save_image_path itself handles. + filename_prefix = expand_date_tokens(filename_prefix, datetime.now()) full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path( filename_prefix, folder_paths.get_output_directory(), width, height) diff --git a/tests/test_filenames.py b/tests/test_filenames.py new file mode 100644 index 0000000..aa8275b --- /dev/null +++ b/tests/test_filenames.py @@ -0,0 +1,38 @@ +# Tests for filename %date:FORMAT% token expansion — part of ComfyUI-AI2Go-Utils. GPL-3.0. +from datetime import datetime + +from nodes.civitai_metadata.filenames import expand_date_tokens + +NOW = datetime(2026, 7, 12, 9, 8, 7) # 2026-07-12 09:08:07 + + +def test_date_with_format_and_subfolder(): + assert expand_date_tokens("%date:yyyy-MM-dd%/Krea-Turbo_9-16", NOW) == "2026-07-12/Krea-Turbo_9-16" + + +def test_date_year_only(): + assert expand_date_tokens("%date:yyyy%/x", NOW) == "2026/x" + + +def test_date_full_datetime(): + assert expand_date_tokens("%date:yyyyMMdd_HHmmss%", NOW) == "20260712_090807" + + +def test_bare_date_defaults_to_iso(): + assert expand_date_tokens("%date%/x", NOW) == "2026-07-12/x" + + +def test_month_vs_minute_case_sensitive(): + # MM = month (07), mm = minute (08) + assert expand_date_tokens("%date:MM-mm%", NOW) == "07-08" + + +def test_no_token_left_untouched(): + assert expand_date_tokens("plain/name", NOW) == "plain/name" + # %year% is get_save_image_path's job — expand_date_tokens must not touch it. + assert expand_date_tokens("%year%/name", NOW) == "%year%/name" + + +def test_empty_and_none(): + assert expand_date_tokens("", NOW) == "" + assert expand_date_tokens(None, NOW) is None diff --git a/tests/test_tracer_model.py b/tests/test_tracer_model.py index 88db53f..b4e4b87 100644 --- a/tests/test_tracer_model.py +++ b/tests/test_tracer_model.py @@ -43,3 +43,32 @@ def test_model_only_lora_and_passthrough(): r = trace(p, "7") assert [lo["name"] for lo in r.loras] == ["C"] assert r.model_name == "myCkpt" + + +def test_checkpoint_folder_is_checkpoints(): + r = trace(_base(), "7") + assert r.model_folder == "checkpoints" + + +def test_unet_loader_diffusion_model(): + # Flux/Krea-style: model comes from a Load Diffusion Model (UNETLoader), not a checkpoint. + p = _base() + del p["1"]["inputs"]["ckpt_name"] + p["1"] = {"class_type": "UNETLoader", "inputs": {"unet_name": "flux/krea-turbo.safetensors", + "weight_dtype": "default"}} + r = trace(p, "7") + assert r.model_name == "krea-turbo" + assert r.model_file == "flux/krea-turbo.safetensors" + assert r.model_folder == "diffusion_models" + + +def test_unet_loader_with_model_only_lora(): + # sampler <- LoraLoaderModelOnly <- UNETLoader (common Flux LoRA wiring) + p = _base() + p["1"] = {"class_type": "UNETLoader", "inputs": {"unet_name": "krea.safetensors", "weight_dtype": "default"}} + p["8"] = {"class_type": "LoraLoaderModelOnly", "inputs": {"model": ["1", 0], + "lora_name": "styleFlux.safetensors", "strength_model": 0.9}} + p["4"]["inputs"]["model"] = ["8", 0] + r = trace(p, "7") + assert r.model_name == "krea" and r.model_folder == "diffusion_models" + assert [lo["name"] for lo in r.loras] == ["styleFlux"] diff --git a/web/js/save_civitai_metadata.js b/web/js/save_civitai_metadata.js index be479d9..230d82b 100644 --- a/web/js/save_civitai_metadata.js +++ b/web/js/save_civitai_metadata.js @@ -13,7 +13,13 @@ const { app } = window.comfyAPI.app; const NODE_IDS = ["AI2GoSaveCivitaiMetadata", "AI2GoSaveCivitaiMetadataAdvanced"]; const SAMPLER_CLASSES = new Set(["KSampler", "KSamplerAdvanced", "SamplerCustom", "SamplerCustomAdvanced"]); const CLIP_CLASSES = new Set(["CLIPTextEncode"]); -const CHECKPOINT_CLASSES = new Set(["CheckpointLoaderSimple", "CheckpointLoader", "unCLIPCheckpointLoader"]); +// class_type -> widget holding the model filename (mirror of tracer.MODEL_SOURCES; the Python side +// also tracks the folder for hashing, which the preview doesn't need). Covers checkpoints and +// standalone diffusion models (UNETLoader "Load Diffusion Model" — Flux/SD3/Krea/etc.). +const MODEL_SOURCES = { + CheckpointLoaderSimple: "ckpt_name", CheckpointLoader: "ckpt_name", unCLIPCheckpointLoader: "ckpt_name", + UNETLoader: "unet_name", UnetLoaderGGUF: "unet_name", UnetLoaderGGUFAdvanced: "unet_name", +}; const LORA_CLASSES = new Set(["LoraLoader", "LoraLoaderModelOnly"]); const BATCH_CLASS = "AI2GoPromptBatch"; @@ -133,9 +139,9 @@ function traceLive(node) { const name = widget(cur, "lora_name"); if (typeof name === "string") r.loras.push({ name: stem(name), strength: widget(cur, "strength_model") }); cur = originNode(cur, "model"); - } else if (CHECKPOINT_CLASSES.has(cls)) { - const ck = widget(cur, "ckpt_name"); - if (typeof ck === "string") r.model_name = stem(ck); + } else if (MODEL_SOURCES[cls]) { + const name = widget(cur, MODEL_SOURCES[cls]); + if (typeof name === "string") r.model_name = stem(name); break; } else { cur = originNode(cur, "model"); @@ -190,16 +196,43 @@ app.registerExtension({ ensureStyles(); chainCallback(nodeType.prototype, "onNodeCreated", function () { const node = this; - const panel = document.createElement("div"); - panel.className = "ai2go-scm"; - panel.textContent = 'Press "🔎 Test detection" to preview what will be saved.'; + const panel = document.createElement("div"); // widget root; ComfyUI pins its height each frame + const content = document.createElement("div"); // natural-height content we measure + content.className = "ai2go-scm"; + content.textContent = 'Press "🔎 Test detection" to preview what will be saved.'; + panel.appendChild(content); + const btn = node.addWidget("button", "🔎 Test detection", null, () => { - try { renderPreview(panel, traceLive(node)); } - catch (e) { panel.textContent = "Trace error: " + e.message; } - node.setDirtyCanvas?.(true, true); + try { renderPreview(content, traceLive(node)); } + catch (e) { content.textContent = "Trace error: " + e.message; } + fit(); }); btn.serialize = false; - node.addDOMWidget("scm_preview", "info", panel, { serialize: false }); + + const w = node.addDOMWidget("scm_preview", "info", panel, { serialize: false }); + + // Reserve the content's natural height so the node grows to fit its text (measured, like the + // Prompt Batch rows editor) — otherwise ComfyUI gives the DOM widget a fixed height and the + // preview overflows below the node border. Measure `content` (not the pinned `panel`) to avoid + // a grow-by-its-own-padding feedback loop. + function fit() { + const h = Math.max(content.scrollHeight, 8); + w.computeSize = () => [node.size?.[0] || 320, h + 8]; + const want = node.computeSize?.(); + if (want) node.setSize([node.size[0], want[1]]); + node.setDirtyCanvas?.(true, true); + } + w.computeSize = () => [node.size?.[0] || 320, Math.max(content.scrollHeight, 8) + 8]; + + // The element attaches after onNodeCreated, so observe it and re-fit the moment it gains (or + // changes) its natural height — self-corrects with no user interaction. + let lastH = 0; + const ro = new ResizeObserver(() => { + const h = content.scrollHeight; + if (h && h !== lastH) { lastH = h; fit(); } + }); + ro.observe(content); + chainCallback(node, "onRemoved", () => ro.disconnect()); }); }, });