fix: model detection for diffusion models, %date% filenames, preview resize

Real-world testing feedback (PR #9):
- Detect models loaded via UNETLoader ("Load Diffusion Model", unet_name in
  the diffusion_models folder) + GGUF unet loaders, not just checkpoints;
  carry the folder so AutoV2 hashing looks in the right place. Fixes
  "model (?)" on Flux/SD3/Krea workflows. Mirrored in the JS Test button.
- Expand %date:FORMAT% filename tokens (frontend feature that doesn't reach a
  custom node's widget) before get_save_image_path, so e.g.
  "%date:yyyy-MM-dd%/Krea-Turbo" works. New comfy-free filenames module.
- Test-detection preview panel now reserves its natural height and grows the
  node to fit (measured, like the Prompt Batch editor) instead of overflowing.

Advanced node sockets kept as-is per user preference. +10 unit tests (39 passed).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

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"))

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

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)
 

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

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"]

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());
     });
   },
 });