diff --git a/backend/server.py b/backend/server.py index 0ee5615..87f9acd 100644 --- a/backend/server.py +++ b/backend/server.py @@ -181,6 +181,18 @@ def create_app(loop: asyncio.AbstractEventLoop) -> web.Application: return web.Response(text=_dumps({"filename": filename}), content_type="application/json") + async def download_file(request: web.Request) -> web.Response: + """Accept a blob POST and return it with Content-Disposition: attachment.""" + body = await request.read() + filename = request.query.get("filename", "workflow.png") + return web.Response( + body=body, + content_type="application/octet-stream", + headers={ + "Content-Disposition": f'attachment; filename="{filename}"', + }, + ) + async def submit_prompt(request: web.Request) -> web.Response: body = await request.json() prompt = body.get("prompt") @@ -249,6 +261,7 @@ def create_app(loop: asyncio.AbstractEventLoop) -> web.Application: app.router.add_get("/files", list_files) app.router.add_get("/browse", browse_dir) app.router.add_post("/upload", upload_file) + app.router.add_post("/download", download_file) app.router.add_post("/prompt", submit_prompt) app.router.add_get("/ws", websocket_handler) diff --git a/desktop.py b/desktop.py index dd9145b..aca7983 100644 --- a/desktop.py +++ b/desktop.py @@ -1,11 +1,13 @@ from __future__ import annotations import asyncio +import base64 import logging import socket import threading import time import urllib.request +from pathlib import Path import webview from aiohttp import web @@ -43,6 +45,34 @@ class _Api: return result[0] return None + def save_workflow_png(self, data_url: str, default_filename: str = "workflow.png") -> str | None: + """Open a native save dialog, write the PNG bytes, and return the saved path.""" + win = self._window_ref[0] + if win is None: + return None + + result = win.create_file_dialog( + webview.SAVE_DIALOG, + save_filename=default_filename, + file_types=( + "PNG image (*.png)", + "All files (*.*)", + ), + ) + if not result: + return None + + path = Path(result[0] if isinstance(result, (list, tuple)) else result).expanduser() + if path.suffix.lower() != ".png": + path = path.with_suffix(".png") + + _, _, encoded = data_url.partition(",") + if not encoded: + raise ValueError("Invalid data URL payload") + + path.write_bytes(base64.b64decode(encoded)) + return str(path) + def _pick_free_port() -> int: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: diff --git a/frontend/src/App.jsx b/frontend/src/App.jsx index 219ada0..ced7656 100644 --- a/frontend/src/App.jsx +++ b/frontend/src/App.jsx @@ -42,6 +42,40 @@ function getOutputSlot(handleId) { return parseInt(handleId.split('::')[1], 10); } +function blobToDataUrl(blob) { + return new Promise((resolve, reject) => { + const reader = new FileReader(); + reader.onloadend = () => resolve(reader.result); + reader.onerror = () => reject(reader.error || new Error('Failed to read file')); + reader.readAsDataURL(blob); + }); +} + +function serializeWorkflowState(nodes, edges) { + return { + version: 1, + nodes: nodes.map((node) => ({ + id: node.id, + type: node.type || 'custom', + position: node.position, + dragHandle: node.dragHandle || '.drag-handle', + data: { + label: node.data?.label || node.data?.className || 'Node', + className: node.data?.className || '', + widgetValues: node.data?.widgetValues || {}, + }, + })), + edges: edges.map((edge) => ({ + id: edge.id, + source: edge.source, + sourceHandle: edge.sourceHandle, + target: edge.target, + targetHandle: edge.targetHandle, + style: edge.style, + })), + }; +} + // ── Graph serialisation → backend prompt format ─────────────────────── function serializeGraph(nodes, edges) { @@ -442,8 +476,12 @@ function Flow() { const defs = nodeDefsRef.current; const hydrated = loadedNodes.map((n) => ({ ...n, + type: n.type || 'custom', + dragHandle: n.dragHandle || '.drag-handle', data: { ...n.data, + label: n.data?.label || n.data?.className || 'Node', + widgetValues: n.data?.widgetValues || {}, definition: defs[n.data.className] || n.data.definition, previewImage: null, tableRows: null, meshData: null, overlay: null, }, @@ -479,11 +517,7 @@ function Flow() { }); if (!blob) throw new Error('Capture returned empty'); - const currentNodes = allNodes.map((n) => ({ - ...n, - data: { ...n.data, previewImage: null, tableRows: null, meshData: null, overlay: null }, - })); - const workflow = { version: 1, nodes: currentNodes, edges: reactFlow.getEdges() }; + const workflow = serializeWorkflowState(allNodes, reactFlow.getEdges()); return embedWorkflow(blob, workflow); }, [reactFlow]); @@ -492,32 +526,62 @@ function Flow() { try { const finalBlob = await getWorkflowBlob(); - if (window.showSaveFilePicker) { - const handle = await window.showSaveFilePicker({ - suggestedName: 'workflow.png', - types: [{ description: 'PNG Image', accept: { 'image/png': ['.png'] } }], - }); - const writable = await handle.createWritable(); - await writable.write(finalBlob); - await writable.close(); - } else { - // Fallback: programmatic download - const a = document.createElement('a'); - a.href = URL.createObjectURL(finalBlob); - a.download = 'workflow.png'; - document.body.appendChild(a); - a.click(); - document.body.removeChild(a); - URL.revokeObjectURL(a.href); + if (window.pywebview?.api?.save_workflow_png) { + const dataUrl = await blobToDataUrl(finalBlob); + const savedPath = await window.pywebview.api.save_workflow_png(dataUrl, 'workflow.png'); + if (!savedPath) { + setStatus({ text: 'Save cancelled.', level: 'info' }); + return; + } + setStatus({ text: `Workflow saved to ${savedPath}.`, level: 'info' }); + return; } - setStatus({ text: 'Workflow saved.', level: 'info' }); - } catch (err) { - if (err.name === 'AbortError') { - setStatus({ text: 'Save cancelled.', level: 'info' }); - } else { - setStatus({ text: 'Save failed: ' + err.message, level: 'error' }); + if ('showSaveFilePicker' in window) { + try { + const handle = await window.showSaveFilePicker({ + suggestedName: 'workflow.png', + types: [ + { + description: 'PNG image', + accept: { 'image/png': ['.png'] }, + }, + ], + }); + const writable = await handle.createWritable(); + await writable.write(finalBlob); + await writable.close(); + setStatus({ text: 'Workflow saved as workflow.png.', level: 'info' }); + return; + } catch (err) { + if (err?.name === 'AbortError') { + setStatus({ text: 'Save cancelled.', level: 'info' }); + return; + } + throw err; + } } + + // Final fallback: trigger a browser download and tell the user where it went. + const resp = await fetch('/download?filename=workflow.png', { + method: 'POST', + body: finalBlob, + }); + const dlBlob = await resp.blob(); + const a = document.createElement('a'); + a.href = URL.createObjectURL(dlBlob); + a.download = 'workflow.png'; + document.body.appendChild(a); + a.click(); + document.body.removeChild(a); + setTimeout(() => URL.revokeObjectURL(a.href), 1000); + + setStatus({ + text: 'Workflow downloaded as workflow.png to your browser default downloads folder.', + level: 'info', + }); + } catch (err) { + setStatus({ text: 'Save failed: ' + err.message, level: 'error' }); } }, [getWorkflowBlob]); @@ -545,7 +609,8 @@ function Flow() { if (!file) return; try { let data; - if (file.name.endsWith('.png') || file.type === 'image/png') { + const lowerName = file.name.toLowerCase(); + if (lowerName.endsWith('.png') || file.type === 'image/png') { data = await extractWorkflow(file); if (!data) { setStatus({ text: 'No workflow data found in image.', level: 'error' }); @@ -571,7 +636,8 @@ function Flow() { event.preventDefault(); const file = files[0]; - if (file.type !== 'image/png') return; + const lowerName = file.name.toLowerCase(); + if (file.type !== 'image/png' && !lowerName.endsWith('.png')) return; try { const data = await extractWorkflow(file); diff --git a/frontend/src/pngMetadata.js b/frontend/src/pngMetadata.js index 3184332..716d3e7 100644 --- a/frontend/src/pngMetadata.js +++ b/frontend/src/pngMetadata.js @@ -1,9 +1,9 @@ /** - * PNG tEXt chunk utilities for embedding/extracting workflow metadata. + * PNG text chunk utilities for embedding/extracting workflow metadata. * * PNG files are composed of chunks: [4-byte length][4-byte type][data][4-byte CRC]. - * We add a tEXt chunk with key "workflow" containing the JSON-serialised graph, - * inserted just before the IEND chunk. + * We add an iTXt chunk with key "workflow" containing the JSON-serialised graph, + * inserted just before the IEND chunk. We still read legacy tEXt chunks. */ // ── CRC32 (PNG uses CRC-32/ISO 3309) ──────────────────────────────── @@ -43,10 +43,64 @@ function chunkType(data, offset) { ); } +function readUint32(data, offset) { + return new DataView(data.buffer, data.byteOffset + offset, 4).getUint32(0); +} + +function buildChunk(type, payload) { + const encoder = new TextEncoder(); + const typeBytes = encoder.encode(type); + const forCrc = new Uint8Array(4 + payload.length); + forCrc.set(typeBytes, 0); + forCrc.set(payload, 4); + + const chunk = new Uint8Array(12 + payload.length); + const view = new DataView(chunk.buffer); + view.setUint32(0, payload.length); + chunk.set(typeBytes, 4); + chunk.set(payload, 8); + view.setUint32(8 + payload.length, crc32(forCrc)); + return chunk; +} + +function parseTextChunk(type, chunkData) { + const decoder = new TextDecoder(); + const keywordEnd = chunkData.indexOf(0); + if (keywordEnd === -1) return null; + + const keyword = decoder.decode(chunkData.subarray(0, keywordEnd)); + if (keyword !== 'workflow') return null; + + if (type === 'tEXt') { + return JSON.parse(decoder.decode(chunkData.subarray(keywordEnd + 1))); + } + + if (type !== 'iTXt') return null; + + const compressionFlagIdx = keywordEnd + 1; + const compressionMethodIdx = keywordEnd + 2; + if (compressionMethodIdx >= chunkData.length) return null; + + const compressionFlag = chunkData[compressionFlagIdx]; + if (compressionFlag !== 0) { + throw new Error('Compressed PNG workflow metadata is not supported'); + } + + let offset = compressionMethodIdx + 1; + const languageEnd = chunkData.indexOf(0, offset); + if (languageEnd === -1) return null; + + offset = languageEnd + 1; + const translatedEnd = chunkData.indexOf(0, offset); + if (translatedEnd === -1) return null; + + return JSON.parse(decoder.decode(chunkData.subarray(translatedEnd + 1))); +} + // ── Public API ─────────────────────────────────────────────────────── /** - * Embed a workflow object into a PNG blob as a tEXt chunk. + * Embed a workflow object into a PNG blob as an iTXt chunk. * Returns a new Blob with the metadata inserted before IEND. */ export async function embedWorkflow(pngBlob, workflow) { @@ -55,33 +109,22 @@ export async function embedWorkflow(pngBlob, workflow) { const encoder = new TextEncoder(); - // Build tEXt payload: keyword \0 text + // Build iTXt payload: + // keyword \0 compression-flag compression-method language-tag \0 translated-keyword \0 text const key = encoder.encode('workflow'); const val = encoder.encode(JSON.stringify(workflow)); - const payload = new Uint8Array(key.length + 1 + val.length); + const payload = new Uint8Array(key.length + 5 + val.length); payload.set(key, 0); - // payload[key.length] is already 0 (null separator) - payload.set(val, key.length + 1); - - // CRC covers type + payload - const typeBytes = encoder.encode('tEXt'); - const forCrc = new Uint8Array(4 + payload.length); - forCrc.set(typeBytes, 0); - forCrc.set(payload, 4); - - // Assemble chunk: length(4) + type(4) + payload + crc(4) - const chunk = new Uint8Array(12 + payload.length); - const view = new DataView(chunk.buffer); - view.setUint32(0, payload.length); - chunk.set(typeBytes, 4); - chunk.set(payload, 8); - view.setUint32(8 + payload.length, crc32(forCrc)); + payload.set(val, key.length + 5); + const chunk = buildChunk('iTXt', payload); // Locate IEND let pos = 8; let iendPos = data.length; while (pos < data.length) { - const len = new DataView(data.buffer, pos, 4).getUint32(0); + if (pos + 8 > data.length) break; + const len = readUint32(data, pos); + if (pos + 12 + len > data.length) break; if (chunkType(data, pos) === 'IEND') { iendPos = pos; break; } pos += 12 + len; } @@ -96,34 +139,30 @@ export async function embedWorkflow(pngBlob, workflow) { } /** - * Extract the workflow object from a PNG blob's tEXt chunks. + * Extract the workflow object from a PNG blob's iTXt/tEXt chunks. * Returns the parsed object, or null if no "workflow" key is found. */ export async function extractWorkflow(pngBlob) { const data = new Uint8Array(await pngBlob.arrayBuffer()); if (!isPng(data)) return null; - const decoder = new TextDecoder(); let pos = 8; + let found = null; while (pos + 8 <= data.length) { - const len = new DataView(data.buffer, pos, 4).getUint32(0); + const len = readUint32(data, pos); + if (pos + 12 + len > data.length) break; const type = chunkType(data, pos); - if (type === 'tEXt' && pos + 8 + len <= data.length) { + if (type === 'tEXt' || type === 'iTXt') { const chunkData = data.subarray(pos + 8, pos + 8 + len); - const nullIdx = chunkData.indexOf(0); - if (nullIdx !== -1) { - const k = decoder.decode(chunkData.subarray(0, nullIdx)); - if (k === 'workflow') { - return JSON.parse(decoder.decode(chunkData.subarray(nullIdx + 1))); - } - } + const parsed = parseTextChunk(type, chunkData); + if (parsed) found = parsed; } if (type === 'IEND') break; pos += 12 + len; } - return null; + return found; } diff --git a/frontend/vite.config.js b/frontend/vite.config.js index 2182b78..63e6370 100644 --- a/frontend/vite.config.js +++ b/frontend/vite.config.js @@ -10,6 +10,7 @@ export default defineConfig({ '/files': 'http://127.0.0.1:8188', '/browse': 'http://127.0.0.1:8188', '/upload': 'http://127.0.0.1:8188', + '/download': 'http://127.0.0.1:8188', '/prompt': 'http://127.0.0.1:8188', '/ws': { target: 'http://127.0.0.1:8188',