snapshot working
This commit is contained in:
@@ -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)
|
||||
|
||||
|
||||
30
desktop.py
30
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:
|
||||
|
||||
@@ -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,33 +526,63 @@ function Flow() {
|
||||
try {
|
||||
const finalBlob = await getWorkflowBlob();
|
||||
|
||||
if (window.showSaveFilePicker) {
|
||||
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;
|
||||
}
|
||||
|
||||
if ('showSaveFilePicker' in window) {
|
||||
try {
|
||||
const handle = await window.showSaveFilePicker({
|
||||
suggestedName: 'workflow.png',
|
||||
types: [{ description: 'PNG Image', accept: { 'image/png': ['.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
|
||||
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(finalBlob);
|
||||
a.href = URL.createObjectURL(dlBlob);
|
||||
a.download = 'workflow.png';
|
||||
document.body.appendChild(a);
|
||||
a.click();
|
||||
document.body.removeChild(a);
|
||||
URL.revokeObjectURL(a.href);
|
||||
}
|
||||
setTimeout(() => URL.revokeObjectURL(a.href), 1000);
|
||||
|
||||
setStatus({ text: 'Workflow saved.', level: 'info' });
|
||||
setStatus({
|
||||
text: 'Workflow downloaded as workflow.png to your browser default downloads folder.',
|
||||
level: 'info',
|
||||
});
|
||||
} catch (err) {
|
||||
if (err.name === 'AbortError') {
|
||||
setStatus({ text: 'Save cancelled.', level: 'info' });
|
||||
} else {
|
||||
setStatus({ text: 'Save failed: ' + err.message, level: 'error' });
|
||||
}
|
||||
}
|
||||
}, [getWorkflowBlob]);
|
||||
|
||||
const copySnapshot = useCallback(() => {
|
||||
@@ -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);
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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',
|
||||
|
||||
Reference in New Issue
Block a user