From 1f9b05cd4bacced8402430a22f1cce1259e86dca Mon Sep 17 00:00:00 2001 From: matei jordache Date: Thu, 2 Apr 2026 00:03:44 -0700 Subject: [PATCH] fix folder and file save --- backend/execution.py | 2 + backend/execution_context.py | 7 ++ backend/nodes/save.py | 51 +++++---------- backend/nodes/save_layers.py | 18 ++++-- backend/server.py | 20 ++++++ frontend/src/App.tsx | 122 +++++++++++++++++++++++++++++++---- frontend/src/nativePicker.ts | 18 +++++- frontend/vite.config.js | 1 + 8 files changed, 183 insertions(+), 56 deletions(-) diff --git a/backend/execution.py b/backend/execution.py index 519de62..b9bdef1 100644 --- a/backend/execution.py +++ b/backend/execution.py @@ -70,6 +70,7 @@ class ExecutionEngine: on_overlay: Callable[[str, str], None] | None = None, on_value: Callable[[str, Any], None] | None = None, on_warning: Callable[[str, str], None] | None = None, + on_file_download: Callable[[str, str], None] | None = None, ) -> dict[str, tuple]: """ Execute the workflow described by `prompt`. @@ -100,6 +101,7 @@ class ExecutionEngine: overlay=on_overlay, value=on_value, warning=on_warning, + file_download=on_file_download, ): for node_id in order: node_def = prompt[node_id] diff --git a/backend/execution_context.py b/backend/execution_context.py index 40020cb..1c96725 100644 --- a/backend/execution_context.py +++ b/backend/execution_context.py @@ -20,6 +20,7 @@ _LEGACY_CALLBACK_ATTRS = { "overlay": "_broadcast_overlay_fn", "value": "_broadcast_value_fn", "warning": "_broadcast_warning_fn", + "file_download": "_broadcast_file_download_fn", } @@ -32,6 +33,7 @@ def execution_callbacks( overlay: Callback | None = None, value: Callback | None = None, warning: Callback | None = None, + file_download: Callback | None = None, ): token = _callbacks_var.set({ "preview": preview, @@ -40,6 +42,7 @@ def execution_callbacks( "overlay": overlay, "value": value, "warning": warning, + "file_download": file_download, }) try: yield @@ -120,3 +123,7 @@ def emit_value(payload: Any) -> None: def emit_warning(message: str) -> None: _emit("warning", message) + + +def emit_file_download(path: str) -> None: + _emit("file_download", path) diff --git a/backend/nodes/save.py b/backend/nodes/save.py index a5f8969..a6f87f4 100644 --- a/backend/nodes/save.py +++ b/backend/nodes/save.py @@ -6,10 +6,14 @@ from pathlib import Path import numpy as np +import tempfile + from backend.node_registry import register_node -from backend.execution_context import emit_warning +from backend.execution_context import emit_warning, emit_file_download from backend.data_types import DataField, LineData, MeshModel, datafield_to_uint8, image_to_uint8 +DOWNLOAD_DIR = Path(tempfile.gettempdir()) / "tono-downloads" + @register_node(display_name="Save") class Save: @classmethod @@ -21,13 +25,6 @@ class Save: "placeholder": "filename", "placement": "top", }), - "directory_path": ("FOLDER_PICKER", { - "default": "", - "label": "directory", - "placement": "top", - "hide_when_input_connected": "directory", - "top_socket_input": "directory", - }), "value": ("DATA_FIELD", { "label": "value", "accepted_types": [ @@ -56,11 +53,11 @@ class Save: }), }, "optional": { - "directory": ("DIRECTORY", {"label": "directory"}), "plot_title": ("STRING", { "default": "", "placeholder": "plot title (optional)", "label": "title", + "show_when_source_type": {"value": ["LINE"]}, }), }, } @@ -80,13 +77,11 @@ class Save: def save( self, filename: str, - directory_path: str, format: str, value, - directory: str | None = None, plot_title: str = "", ): - path = self._resolve_save_path(filename, format, directory, directory_path) + path = self._resolve_save_path(filename, format) if isinstance(value, MeshModel): self._save_mesh(path, value, format) @@ -107,15 +102,10 @@ class Save: raise ValueError(f"Save does not support input type: {type(value).__name__}") self._send_warning(f"Saved to {path.name}") + emit_file_download(str(path)) return () - def _resolve_save_path( - self, - filename: str, - format_name: str, - directory: str | None, - directory_path: str = "", - ) -> Path: + def _resolve_save_path(self, filename: str, format_name: str) -> Path: ext_map = { "PNG": ".png", "TIFF": ".tiff", @@ -129,25 +119,16 @@ class Save: ext = ext_map[format_name] raw_filename = str(filename).strip() if filename is not None else "" - raw_directory = str(directory).strip() if directory is not None else "" - if not raw_directory: - raw_directory = str(directory_path).strip() if directory_path is not None else "" - if not raw_filename: raise ValueError("No output filename selected — enter a file name.") - if raw_directory: - dir_path = Path(raw_directory).expanduser() - if dir_path.exists() and not dir_path.is_dir(): - raise ValueError("Directory input expects a folder path, not a file path.") - if not dir_path.exists(): - if dir_path.suffix: - raise ValueError("Directory input expects a folder path, not a file path.") - dir_path.mkdir(parents=True, exist_ok=True) - path = dir_path / Path(raw_filename).name + candidate = Path(raw_filename).expanduser() + if candidate.is_absolute(): + candidate.parent.mkdir(parents=True, exist_ok=True) + path = candidate else: - path = Path(raw_filename).expanduser() - path.parent.mkdir(parents=True, exist_ok=True) + DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True) + path = DOWNLOAD_DIR / candidate.name if path.suffix.lower() != ext: path = path.with_suffix(ext) @@ -156,7 +137,7 @@ class Save: def _save_datafield(self, path: Path, field: DataField, format_name: str): if format_name == "TIFF": import tifffile - tifffile.imwrite(str(path), np.asarray(field.data, dtype=np.float32)) + tifffile.imwrite(str(path), datafield_to_uint8(field, field.colormap)) return if format_name == "NPZ": np.savez(str(path), field=np.asarray(field.data)) diff --git a/backend/nodes/save_layers.py b/backend/nodes/save_layers.py index 14c832c..022c883 100644 --- a/backend/nodes/save_layers.py +++ b/backend/nodes/save_layers.py @@ -4,7 +4,7 @@ import numpy as np from pathlib import Path from backend.node_registry import register_node -from backend.execution_context import emit_warning +from backend.execution_context import emit_warning, emit_file_download from backend.data_types import DataField, image_to_uint8 from backend.nodes.helpers import _MAX_SAVE_FIELDS @@ -35,9 +35,10 @@ class SaveImage: "placeholder": "filename", "placement": "top", }), - "directory_path": ("FOLDER_PICKER", { + "directory_path": ("STRING", { "default": "", "label": "directory", + "placeholder": "directory (optional, desktop only)", "placement": "top", "hide_when_input_connected": "directory", "top_socket_input": "directory", @@ -92,6 +93,7 @@ class SaveImage: self._save_npz(path, layers, layer_names) self._send_warning(f"Saved {len(layers)} layer(s) to {path.name}") + emit_file_download(str(path)) return () def _save_tiff(self, path: Path, layers: list[DataField | np.ndarray], layer_names: list[str]): @@ -140,9 +142,15 @@ class SaveImage: path = dir_path / filename_part else: if not raw_filename: - raise ValueError("No output path selected — use Browse to pick a location.") - path = Path(raw_filename).expanduser() - path.parent.mkdir(parents=True, exist_ok=True) + raise ValueError("No output filename selected — enter a file name.") + candidate = Path(raw_filename).expanduser() + if candidate.is_absolute(): + candidate.parent.mkdir(parents=True, exist_ok=True) + path = candidate + else: + from backend.nodes.save import DOWNLOAD_DIR + DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True) + path = DOWNLOAD_DIR / candidate.name if path.suffix.lower() != ext: path = path.with_suffix(ext) diff --git a/backend/server.py b/backend/server.py index 0887561..8c2f909 100644 --- a/backend/server.py +++ b/backend/server.py @@ -32,6 +32,7 @@ import asyncio import json import logging import math +import secrets import sys from collections import defaultdict from copy import deepcopy @@ -139,6 +140,7 @@ def create_app( session_engines: dict[str, ExecutionEngine] = {} session_websockets: dict[str, set[web.WebSocketResponse]] = defaultdict(set) + pending_downloads: dict[str, Path] = {} def _is_link(value) -> bool: return ( @@ -254,6 +256,12 @@ def create_app( def on_warning(session_id: str, node_id: str, message: str) -> None: broadcast(session_id, {"type": "node_warning", "data": {"node_id": node_id, "message": message}}) + def on_file_download(session_id: str, node_id: str, file_path: str) -> None: + token = secrets.token_urlsafe(16) + path = Path(file_path) + pending_downloads[token] = path + broadcast(session_id, {"type": "file_download", "data": {"node_id": node_id, "token": token, "filename": path.name}}) + async def index(request: web.Request) -> web.Response: if not getattr(sys, "frozen", False): try: @@ -470,6 +478,16 @@ def create_app( headers={"Content-Disposition": f'attachment; filename="{filename}"'}, ) + async def download_saved_file(request: web.Request) -> web.Response: + token = request.match_info["token"] + path = pending_downloads.pop(token, None) + if path is None or not path.is_file(): + raise web.HTTPNotFound(reason="File not found") + return web.FileResponse( + path, + headers={"Content-Disposition": f'attachment; filename="{path.name}"'}, + ) + async def save_workflow_png(request: web.Request) -> web.Response: body = await request.read() target_path = request.query.get("path", "") @@ -535,6 +553,7 @@ def create_app( on_overlay=lambda node_id, overlay_data: on_overlay(session_id, node_id, overlay_data), on_value=lambda node_id, payload: on_value(session_id, node_id, payload), on_warning=lambda node_id, message: on_warning(session_id, node_id, message), + on_file_download=lambda node_id, file_path: on_file_download(session_id, node_id, file_path), ), ) broadcast(session_id, {"type": "execution_complete", "data": {"prompt_id": prompt_id}}) @@ -627,6 +646,7 @@ def create_app( app.router.add_get("/help-docs", get_help_docs) app.router.add_get("/help-docs/{filename}", get_help_doc_file) app.router.add_post("/prompt", submit_prompt) + app.router.add_get("/download-save/{token}", download_saved_file) app.router.add_get("/check-update", check_update) app.router.add_get("/ws", websocket_handler) diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx index d7feb8b..5170003 100644 --- a/frontend/src/App.tsx +++ b/frontend/src/App.tsx @@ -906,6 +906,7 @@ function Flow() { const nextIdRef = useRef(1); const autoRunTimer = useRef | null>(null); const autoRunRef = useRef<(() => void) | null>(null); + const pendingBrowserFilesRef = useRef>(new Map()); const defaultWorkflowLoadAttemptedRef = useRef(false); const lastPastedClipboardTextRef = useRef(''); const pasteRepeatCountRef = useRef(0); @@ -1305,7 +1306,34 @@ function Flow() { }, [getResolvedPathInput, reactFlow, setNodeOutputs]); const refreshFolderNodeOutputs = useCallback(async (nodeId: string, folderPath: any) => { - const entries = folderPath ? await api.getFolderFiles(folderPath) : []; + let entries: any[] = []; + + if (folderPath) { + // Check for pending browser files first (folder was picked but files not yet uploaded) + const prefix = String(folderPath).endsWith('/') ? String(folderPath) : String(folderPath) + '/'; + const pendingEntries: any[] = []; + for (const uri of pendingBrowserFilesRef.current.keys()) { + if (uri.startsWith(prefix)) { + const name = uri.slice(prefix.length); + // Skip files in subdirectories for the top-level listing + if (!name.includes('/')) { + pendingEntries.push({ name, type: 'FILE_PATH', path: uri }); + } + } + } + + if (pendingEntries.length > 0) { + // Build listing locally from pending files + entries = [ + { name: 'directory', type: 'DIRECTORY', path: folderPath }, + ...pendingEntries.sort((a: any, b: any) => a.name.localeCompare(b.name)), + ]; + } else { + // Fall back to server (native builds, or files already uploaded) + entries = await api.getFolderFiles(folderPath); + } + } + setNodeOutputs( nodeId, entries.map((entry: any) => entry.type), @@ -1407,6 +1435,24 @@ function Flow() { case 'node_warning': updateNodeData(msg.data.node_id, { warning: msg.data.message }); break; + case 'file_download': { + const dlToken = msg.data.token; + const dlFilename = msg.data.filename || 'download'; + fetch(`/download-save/${encodeURIComponent(dlToken)}`) + .then((r) => r.ok ? r.blob() : Promise.reject(new Error(`Download failed: ${r.status}`))) + .then((blob) => { + const url = URL.createObjectURL(blob); + const a = document.createElement('a'); + a.href = url; + a.download = dlFilename; + document.body.appendChild(a); + a.click(); + a.remove(); + URL.revokeObjectURL(url); + }) + .catch((err) => setStatus({ text: String(err.message), level: 'error' })); + break; + } case 'nodes_updated': api.getNodes().then((defs) => { nodeDefsRef.current = defs; @@ -1675,16 +1721,21 @@ function Flow() { throw new Error('Selected folder is empty or could not be read.'); } + const folder = await api.createUploadFolder(rootName); + const folderUri = folder.path; // e.g. "session://uploads/myfolder" + + // Store File objects for lazy upload — only uploaded when actually used + for (const entry of selection.entries || []) { + const fileUri = `session://uploads/${entry.relativePath}`; + pendingBrowserFilesRef.current.set(fileUri, entry.file); + } + setStatus({ - text: `Importing folder "${rootName}" into this session…`, + text: `Folder "${rootName}" loaded (${(selection.entries || []).length} files).`, level: 'info', }); - const folder = await api.createUploadFolder(rootName); - for (const entry of selection.entries || []) { - await api.uploadFile(entry.file, { relativePath: entry.relativePath }); - } - return folder.path; + return folderUri; } const [entry] = selection.entries || []; @@ -1729,19 +1780,63 @@ function Flow() { } }, [uploadBrowserSelection]); + // ── Lazy upload of pending browser files ───────────────────────────── + + const uploadPendingFiles = useCallback(async (prompt: Record) => { + const pending = pendingBrowserFilesRef.current; + if (pending.size === 0) return; + + // Collect all string values from the prompt (folder paths and file paths) + const promptValues = new Set(); + for (const nodeData of Object.values(prompt)) { + const inputs = (nodeData as any)?.inputs; + if (!inputs || typeof inputs !== 'object') continue; + for (const val of Object.values(inputs)) { + if (typeof val === 'string') promptValues.add(val); + } + } + + // Upload files that are directly referenced OR inside a referenced folder + const toUpload = new Set(); + for (const uri of pending.keys()) { + if (promptValues.has(uri)) { + toUpload.add(uri); + continue; + } + // Check if any prompt value is a folder prefix of this file + for (const pv of promptValues) { + const prefix = pv.endsWith('/') ? pv : pv + '/'; + if (uri.startsWith(prefix)) { + toUpload.add(uri); + break; + } + } + } + + for (const uri of toUpload) { + const file = pending.get(uri)!; + const relativePath = uri.replace(/^session:\/\/uploads\//, ''); + await api.uploadFile(file, { relativePath }); + pending.delete(uri); + } + }, []); + // ── Node context value (stable) ───────────────────────────────────── - const onManualTrigger = useCallback((nodeId: string) => { + const onManualTrigger = useCallback(async (nodeId: string) => { const currentNodes = (reactFlow.getNodes() as TonoNode[]); const currentEdges = (reactFlow.getEdges() as TonoEdge[]); // Include ALL nodes (no excludeManualTrigger) so the save node is in the prompt const prompt = serializeExecutionGraph(currentNodes, currentEdges); if (!prompt || Object.keys(prompt).length === 0) return; setStatus({ text: 'Saving…', level: 'info' }); - api.runPrompt(prompt).catch((err) => { + try { + await uploadPendingFiles(prompt); + await api.runPrompt(prompt); + } catch (err: any) { setStatus({ text: 'Save failed: ' + err.message, level: 'error' }); - }); - }, [reactFlow]); + } + }, [reactFlow, uploadPendingFiles]); const openJournalTab = useCallback(() => { setHelpTabs((prev) => { @@ -1874,11 +1969,12 @@ function Flow() { } setStatus({ text: 'Running…', level: 'info' }); try { + await uploadPendingFiles(prompt); await api.runPrompt(prompt); } catch (err: any) { setStatus({ text: 'Failed: ' + err.message, level: 'error' }); } - }, [reactFlow]); + }, [reactFlow, uploadPendingFiles]); // Debounced auto-run via ref to avoid dependency chains autoRunRef.current = () => { @@ -1895,7 +1991,7 @@ function Flow() { const prompt = serializeExecutionGraph(currentNodes, currentEdges, { excludeManualTrigger: true }); if (!prompt || Object.keys(prompt).length === 0) return; setStatus({ text: 'Running…', level: 'info' }); - api.runPrompt(prompt).catch((err) => { + uploadPendingFiles(prompt).then(() => api.runPrompt(prompt)).catch((err) => { setStatus({ text: 'Failed: ' + err.message, level: 'error' }); }); }; diff --git a/frontend/src/nativePicker.ts b/frontend/src/nativePicker.ts index fb4c579..89b316c 100644 --- a/frontend/src/nativePicker.ts +++ b/frontend/src/nativePicker.ts @@ -1,9 +1,18 @@ -const FILE_ACCEPT = [ +const SUPPORTED_EXTENSIONS = new Set([ '.png', '.jpg', '.jpeg', '.tiff', '.tif', '.bmp', '.npy', '.npz', '.gwy', '.sxm', '.ibw', + '.h5', '.hdf5', '.ttf', '.otf', '.woff', '.woff2', -].join(','); +]); + +const FILE_ACCEPT = [...SUPPORTED_EXTENSIONS].join(','); + +function hasSupportedExtension(name: string): boolean { + const dot = name.lastIndexOf('.'); + if (dot < 0) return false; + return SUPPORTED_EXTENSIONS.has(name.slice(dot).toLowerCase()); +} function normalizeRelativePath(path: string) { return String(path || '').replace(/\\/g, '/').replace(/^\/+/, ''); @@ -48,6 +57,7 @@ async function collectDirectoryEntries(handle: FileSystemDirectoryHandle, prefix for await (const [name, child] of (handle as any).entries()) { const relativePath = prefix ? `${prefix}/${name}` : name; if (child.kind === 'file') { + if (!hasSupportedExtension(name)) continue; const file = await child.getFile(); entries.push({ file, relativePath: normalizeRelativePath(relativePath) }); continue; @@ -108,7 +118,9 @@ export async function pickNativeDirectorySelection() { return null; } - const files = await pickWithInput({ directory: true }); + const allFiles = await pickWithInput({ directory: true }); + if (allFiles.length === 0) return null; + const files = allFiles.filter((f) => hasSupportedExtension(f.name)); if (files.length === 0) return null; const entries = files.map((file: File) => ({ file, diff --git a/frontend/vite.config.js b/frontend/vite.config.js index bacb73c..a5fc152 100644 --- a/frontend/vite.config.js +++ b/frontend/vite.config.js @@ -18,6 +18,7 @@ export default defineConfig({ '/channels': 'http://127.0.0.1:8188', '/upload-folder': 'http://127.0.0.1:8188', '/upload': 'http://127.0.0.1:8188', + '/download-save': 'http://127.0.0.1:8188', '/download': 'http://127.0.0.1:8188', '/file-content': 'http://127.0.0.1:8188', '/help-docs': { target: 'http://127.0.0.1:8188', changeOrigin: true },