snapshot working

This commit is contained in:
2026-03-23 22:31:49 -07:00
parent de2d90187a
commit 29107bc141
5 changed files with 214 additions and 65 deletions

View File

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

View File

@@ -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:

View File

@@ -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);

View 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;
}

View File

@@ -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',