diff --git a/__init__.py b/__init__.py index 0ca6551..32269fe 100644 --- a/__init__.py +++ b/__init__.py @@ -1,6 +1,7 @@ from .dumas_image_nodes import ( NODE_CLASS_MAPPINGS as IMAGE_NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as IMAGE_NODE_DISPLAY_NAME_MAPPINGS, + _next_counter, resolve_serve_token, ) from .dumas_json_nodes import ( @@ -84,6 +85,35 @@ try: return web.Response(status=404, text="unknown or expired preview token") return web.FileResponse(path) + @PromptServer.instance.routes.get("/dumas/api/save_image/next_counter") + async def dumas_save_image_next_counter(request): + folder_raw = str(request.query.get("folder", "") or "").strip() + name = str(request.query.get("name", "") or "") + try: + digits = max(1, min(8, int(request.query.get("digits", "3")))) + except Exception: + digits = 3 + + def _scan(): + import folder_paths + + output_dir = folder_paths.get_output_directory() + base = os.path.abspath(folder_raw) if folder_raw else output_dir + parts = [part for part in name.replace("\\", "/").split("/") if part] + if not parts: + return 1, "" + frame_dir = os.path.join(base, *parts[:-1]) + counter = _next_counter(frame_dir, parts[-1]) + resolved_name = parts[-1].replace("%counter%", str(counter).zfill(digits)) + return counter, "/".join(parts[:-1] + [resolved_name]) + + try: + loop = asyncio.get_running_loop() + counter, resolved = await loop.run_in_executor(None, _scan) + return web.json_response({"ok": True, "counter": counter, "resolved": resolved}) + except Exception as exc: + return web.json_response({"ok": False, "message": str(exc), "counter": 1, "resolved": ""}) + @PromptServer.instance.routes.get("/dumas/api/pick_directory") async def dumas_pick_directory(request): try: diff --git a/js/save_image/index.js b/js/save_image/index.js index 6661c8c..a81c139 100644 --- a/js/save_image/index.js +++ b/js/save_image/index.js @@ -15,6 +15,7 @@ const CHIPS = [ { label: "+ Height", token: "%height%" }, { label: "+ Batch #", token: "%batch_num%" }, ]; +const COUNTER_DIGITS = 3; let cssDone = false; @@ -127,11 +128,47 @@ function buildPreviewPath(node, ui) { .replaceAll("%input2%", "input2") .replaceAll("%width%", String(node._dsiImageWidth || 0)) .replaceAll("%height%", String(node._dsiImageHeight || 0)) - .replaceAll("%batch_num%", "0") - .replaceAll("%counter%", "001"); + .replaceAll("%batch_num%", "0"); + if ( + ui._dsiCounterResolved + && ui._dsiCounterKey === `${folder}\u0000${pattern}${format}` + ) { + return `${folder}/${ui._dsiCounterResolved}`.replaceAll("\\", "/"); + } + pattern = pattern.replaceAll("%counter%", String(ui._dsiNextCounter || 1).padStart(COUNTER_DIGITS, "0")); return `${folder}/${pattern}${format}`.replaceAll("\\", "/"); } +function scheduleCounterFetch(node, ui) { + const folder = (ui.folderInput.value || "").trim(); + const format = ui.format === "jpg" ? ".jpg" : ".png"; + let pattern = ui.patternInput.value || DEFAULT_PATTERN; + pattern = resolveDateTokens(pattern) + .replaceAll("%input%", "input") + .replaceAll("%input2%", "input2") + .replaceAll("%width%", String(node._dsiImageWidth || 0)) + .replaceAll("%height%", String(node._dsiImageHeight || 0)) + .replaceAll("%batch_num%", "0"); + const key = `${folder}\u0000${pattern}${format}`; + if (ui._dsiCounterKey === key && ui._dsiCounterResolved) return; + clearTimeout(ui._dsiCounterTimer); + ui._dsiCounterTimer = setTimeout(async () => { + try { + const response = await fetch( + `/dumas/api/save_image/next_counter?folder=${encodeURIComponent(folder)}&name=${encodeURIComponent(pattern + format)}&digits=${COUNTER_DIGITS}`, + ); + const data = await response.json(); + if (!data?.ok) return; + ui._dsiNextCounter = data.counter || 1; + ui._dsiCounterResolved = data.resolved || ""; + ui._dsiCounterKey = key; + ui.prevPath.textContent = buildPreviewPath(node, ui); + } catch (_error) { + // Leave the fallback preview in place. + } + }, 150); +} + function imageUrl(meta) { return `/view?filename=${encodeURIComponent(meta.filename)}&type=${encodeURIComponent(meta.type)}&subfolder=${encodeURIComponent(meta.subfolder || "")}&t=${Date.now()}`; } @@ -175,6 +212,7 @@ function renderPreview(node, ui) { const currentIndex = Math.max(0, Math.min(node._dsiImageIndex || 0, images.length - 1)); node._dsiImageIndex = currentIndex; ui.prevPath.textContent = buildPreviewPath(node, ui); + scheduleCounterFetch(node, ui); ui.folderButton.title = (ui.folderInput.value || "").trim() ? `Copy folder path: ${(ui.folderInput.value || "").trim()}` : "Using ComfyUI output folder"; @@ -339,6 +377,10 @@ function createRoot(node) { format: "png", saveOnRun: true, embedWorkflow: true, + _dsiNextCounter: 1, + _dsiCounterResolved: "", + _dsiCounterKey: "", + _dsiCounterTimer: null, }; folderInput.addEventListener("input", () => { diff --git a/tests/test_dumas_image_nodes.py b/tests/test_dumas_image_nodes.py index f3b2e82..c4bada2 100644 --- a/tests/test_dumas_image_nodes.py +++ b/tests/test_dumas_image_nodes.py @@ -2,6 +2,7 @@ import importlib import os import sys import tempfile +import time import types import unittest from unittest import mock @@ -151,6 +152,58 @@ class DumasImageNodeTests(unittest.TestCase): saved_name = os.path.basename(FakePILImage.saved_paths[0][0]) self.assertEqual(saved_name, "shot_alpha_beta_final_001.png") + def test_save_image_resolves_date_size_and_batch_tokens(self): + node = self.image_nodes.DumasSaveImageNode() + image = FakeTensorBatch(width=10, height=12, count=2) + fake_now = time.struct_time((2026, 8, 5, 13, 7, 9, 2, 217, -1)) + + with mock.patch.object(self.image_nodes.time, "localtime", return_value=fake_now): + node.save_images( + images=image, + folder=self.temp_dir, + pattern="asset_%date:yyyy-MM-dd%_%date:hh-mm-ss%_%width%x%height%_%batch_num%_%counter%", + format="png", + quality=100, + embed_workflow=False, + save_on_run=True, + ) + + saved_names = [os.path.basename(path) for path, _args, _kwargs in FakePILImage.saved_paths] + self.assertEqual( + saved_names, + [ + "asset_2026-08-05_13-07-09_10x12_0_001.png", + "asset_2026-08-05_13-07-09_10x12_1_001.png", + ], + ) + + def test_save_image_counter_increments_for_existing_files(self): + node = self.image_nodes.DumasSaveImageNode() + image = FakeTensorBatch() + + node.save_images( + images=image, + folder=self.temp_dir, + pattern="counter_%counter%", + format="png", + quality=100, + embed_workflow=False, + save_on_run=True, + ) + open(FakePILImage.saved_paths[0][0], "a", encoding="utf-8").close() + node.save_images( + images=image, + folder=self.temp_dir, + pattern="counter_%counter%", + format="png", + quality=100, + embed_workflow=False, + save_on_run=True, + ) + + saved_names = [os.path.basename(path) for path, _args, _kwargs in FakePILImage.saved_paths] + self.assertEqual(saved_names, ["counter_001.png", "counter_002.png"]) + def test_save_image_returns_ui_entries_for_output_folder(self): node = self.image_nodes.DumasSaveImageNode() image = FakeTensorBatch()