From 8191dabb6140f613ef5e1e6ba84dfd820bad286f Mon Sep 17 00:00:00 2001 From: Chris Dumas Date: Wed, 5 Aug 2026 13:15:39 +0000 Subject: [PATCH] Add save image browse and external preview --- __init__.py | 34 +++++++++++++++++++++++++++++ dumas_image_nodes.py | 17 +++++++++++++++ js/save_image/index.js | 38 ++++++++++++++++++++++++++++----- tests/test_dumas_image_nodes.py | 2 ++ 4 files changed, 86 insertions(+), 5 deletions(-) diff --git a/__init__.py b/__init__.py index 5cd4c80..6a201bc 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, + resolve_serve_token, ) from .dumas_json_nodes import ( NODE_CLASS_MAPPINGS as JSON_NODE_CLASS_MAPPINGS, @@ -17,4 +18,37 @@ NODE_DISPLAY_NAME_MAPPINGS.update(IMAGE_NODE_DISPLAY_NAME_MAPPINGS) WEB_DIRECTORY = "./js" +try: + import os + from aiohttp import web + from server import PromptServer + + @PromptServer.instance.routes.get("/dumas/api/save_image/file") + async def dumas_save_image_file(request): + path = resolve_serve_token(request.query.get("t", "")) + if not path or not os.path.isfile(path): + return web.Response(status=404, text="unknown or expired preview token") + return web.FileResponse(path) + + @PromptServer.instance.routes.get("/dumas/api/pick_directory") + async def dumas_pick_directory(_request): + try: + import tkinter as tk + from tkinter import filedialog + + root = tk.Tk() + root.withdraw() + root.attributes("-topmost", True) + selected = filedialog.askdirectory( + initialdir=str(_request.query.get("path", "") or "") or None + ) + root.destroy() + if not selected: + return web.json_response({"ok": False, "cancelled": True}) + return web.json_response({"ok": True, "path": selected}) + except Exception as exc: + return web.json_response({"ok": False, "message": str(exc)}) +except Exception: + pass + __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] diff --git a/dumas_image_nodes.py b/dumas_image_nodes.py index 52b8ce6..99fa1f2 100644 --- a/dumas_image_nodes.py +++ b/dumas_image_nodes.py @@ -3,6 +3,8 @@ import os import random import re import time +import uuid +from collections import OrderedDict import numpy as np from PIL import Image @@ -14,6 +16,8 @@ _MEDIA_EXT_RE = re.compile( r"\.(png|jpe?g|webp|gif|bmp|tiff?|avif|mp4|mov|webm|mkv|m4v)$", re.IGNORECASE ) _DATE_TOKEN_RE = re.compile(r"%date:([^%]+)%") +_SERVE_TOKENS = OrderedDict() +_SERVE_CAP = 256 def _clean_input_token_value(value): @@ -95,6 +99,18 @@ def _safe_pattern(value): return value +def _register_serve_token(path): + token = uuid.uuid4().hex + _SERVE_TOKENS[token] = path + while len(_SERVE_TOKENS) > _SERVE_CAP: + _SERVE_TOKENS.popitem(last=False) + return token + + +def resolve_serve_token(token): + return _SERVE_TOKENS.get(str(token or "")) + + def _is_within_directory(parent_path, child_path): try: return os.path.commonpath([parent_path, child_path]) == parent_path @@ -378,6 +394,7 @@ class DumasSaveImageNode: "filename": filename, "subfolder": frame_dir.replace("\\", "/"), "type": "external", + "token": _register_serve_token(full_path), } ) diff --git a/js/save_image/index.js b/js/save_image/index.js index 57551e6..6661c8c 100644 --- a/js/save_image/index.js +++ b/js/save_image/index.js @@ -136,8 +136,15 @@ function imageUrl(meta) { return `/view?filename=${encodeURIComponent(meta.filename)}&type=${encodeURIComponent(meta.type)}&subfolder=${encodeURIComponent(meta.subfolder || "")}&t=${Date.now()}`; } +function entrySrc(meta) { + if (!meta) return null; + if (meta.type === "output" && meta.filename) return imageUrl(meta); + if (meta.token) return `/dumas/api/save_image/file?t=${encodeURIComponent(meta.token)}`; + return null; +} + function canPreview(meta) { - return meta?.type === "output"; + return !!entrySrc(meta); } function updateActionState(node, ui) { @@ -154,7 +161,9 @@ function updateActionState(node, ui) { async function copyCurrentImage(node) { const entry = node._dsiImages?.[node._dsiImageIndex || 0]; if (!entry || !navigator.clipboard?.write || typeof ClipboardItem === "undefined") return false; - const response = await fetch(imageUrl(entry)); + const src = entrySrc(entry); + if (!src) return false; + const response = await fetch(src); if (!response.ok) return false; const blob = await response.blob(); await navigator.clipboard.write([new ClipboardItem({ [blob.type || "image/png"]: blob })]); @@ -188,7 +197,7 @@ function renderPreview(node, ui) { ui.view.classList.add("has"); ui.placeholder.style.display = "none"; ui.img.style.display = "block"; - ui.img.src = imageUrl(current); + ui.img.src = entrySrc(current); ui.info.textContent = images.length > 1 ? `Showing saved image ${currentIndex + 1} of ${images.length}` : "Showing latest saved image"; @@ -341,7 +350,24 @@ function createRoot(node) { renderPreview(node, ui); }); browseButton.addEventListener("click", async () => { - ui.info.textContent = "Browser folder picking cannot provide a real filesystem path here. Paste the full path into the folder field."; + try { + const response = await fetch(`/dumas/api/pick_directory?path=${encodeURIComponent((folderInput.value || "").trim())}`); + const data = await response.json(); + if (data?.ok && data.path) { + folderInput.value = data.path; + setWidgetValue(node, "folder", folderInput.value); + renderPreview(node, ui); + ui.info.textContent = `Selected folder: ${data.path}`; + return; + } + if (data?.cancelled) { + ui.info.textContent = "Folder selection cancelled"; + return; + } + ui.info.textContent = data?.message || "Folder picker is unavailable here"; + } catch (_error) { + ui.info.textContent = "Folder picker is unavailable here"; + } }); fmtPng.addEventListener("click", () => { ui.format = "png"; @@ -366,7 +392,9 @@ function createRoot(node) { openButton.addEventListener("click", () => { const entry = node._dsiImages?.[node._dsiImageIndex || 0]; if (!entry) return; - window.open(imageUrl(entry), "_blank", "noopener,noreferrer"); + const src = entrySrc(entry); + if (!src) return; + window.open(src, "_blank", "noopener,noreferrer"); }); copyButton.addEventListener("click", async () => { try { diff --git a/tests/test_dumas_image_nodes.py b/tests/test_dumas_image_nodes.py index 69143a2..f3b2e82 100644 --- a/tests/test_dumas_image_nodes.py +++ b/tests/test_dumas_image_nodes.py @@ -185,6 +185,7 @@ class DumasImageNodeTests(unittest.TestCase): self.assertEqual(result["ui"]["images"][0]["type"], "external") self.assertEqual(result["ui"]["images"][0]["subfolder"], external_dir.replace("\\", "/")) + self.assertTrue(result["ui"]["images"][0]["token"]) def test_save_image_handles_windows_different_drive_paths(self): node = self.image_nodes.DumasSaveImageNode() @@ -211,6 +212,7 @@ class DumasImageNodeTests(unittest.TestCase): self.assertEqual(result["ui"]["images"][0]["type"], "external") self.assertEqual(result["ui"]["images"][0]["subfolder"], "D:/renders") + self.assertTrue(result["ui"]["images"][0]["token"]) def test_save_image_skips_when_save_is_disabled(self): node = self.image_nodes.DumasSaveImageNode()