Add save image browse and external preview

2026-08-05 13:15:39 +00:00
parent 17f7adba65
commit 8191dabb61
4 changed files with 86 additions and 5 deletions
@@ -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"]
@@ -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),
                    }
                )

@@ -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 {
@@ -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()