From 62a371e12b4763bf6f9aeb42ff4928138df6ae26 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Thu, 18 May 2023 02:41:21 -0400
Subject: [PATCH] Load workflow from latent file.

---
 nodes.py               |  2 +-
 web/scripts/app.js     |  7 ++++++-
 web/scripts/pnginfo.js | 16 ++++++++++++++++
 web/scripts/ui.js      |  2 +-
 4 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/nodes.py b/nodes.py
index 7b450df3..3c61cd2e 100644
--- a/nodes.py
+++ b/nodes.py
@@ -274,7 +274,7 @@ class SaveLatent:
         if prompt is not None:
             prompt_info = json.dumps(prompt)
 
-        metadata = {"workflow": prompt_info}
+        metadata = {"prompt": prompt_info}
         if extra_pnginfo is not None:
             for x in extra_pnginfo:
                 metadata[x] = json.dumps(extra_pnginfo[x])
diff --git a/web/scripts/app.js b/web/scripts/app.js
index ef3b44c8..97b7c8d3 100644
--- a/web/scripts/app.js
+++ b/web/scripts/app.js
@@ -2,7 +2,7 @@ import { ComfyWidgets } from "./widgets.js";
 import { ComfyUI, $el } from "./ui.js";
 import { api } from "./api.js";
 import { defaultGraph } from "./defaultGraph.js";
-import { getPngMetadata, importA1111 } from "./pnginfo.js";
+import { getPngMetadata, importA1111, getLatentMetadata } from "./pnginfo.js";
 
 /** 
  * @typedef {import("types/comfy").ComfyExtension} ComfyExtension
@@ -1308,6 +1308,11 @@ export class ComfyApp {
 				this.loadGraphData(JSON.parse(reader.result));
 			};
 			reader.readAsText(file);
+		} else if (file.name?.endsWith(".latent")) {
+			const info = await getLatentMetadata(file);
+			if (info.workflow) {
+				this.loadGraphData(JSON.parse(info.workflow));
+			}
 		}
 	}
 
diff --git a/web/scripts/pnginfo.js b/web/scripts/pnginfo.js
index 209b562a..8ddb7a1c 100644
--- a/web/scripts/pnginfo.js
+++ b/web/scripts/pnginfo.js
@@ -47,6 +47,22 @@ export function getPngMetadata(file) {
 	});
 }
 
+export function getLatentMetadata(file) {
+	return new Promise((r) => {
+		const reader = new FileReader();
+		reader.onload = (event) => {
+			const safetensorsData = new Uint8Array(event.target.result);
+			const dataView = new DataView(safetensorsData.buffer);
+			let header_size = dataView.getUint32(0, true);
+			let offset = 8;
+			let header = JSON.parse(String.fromCharCode(...safetensorsData.slice(offset, offset + header_size)));
+			r(header.__metadata__);
+		};
+
+		reader.readAsArrayBuffer(file);
+	});
+}
+
 export async function importA1111(graph, parameters) {
 	const p = parameters.lastIndexOf("\nSteps:");
 	if (p > -1) {
diff --git a/web/scripts/ui.js b/web/scripts/ui.js
index 77517aec..2c9043d0 100644
--- a/web/scripts/ui.js
+++ b/web/scripts/ui.js
@@ -465,7 +465,7 @@ export class ComfyUI {
 		const fileInput = $el("input", {
 			id: "comfy-file-input",
 			type: "file",
-			accept: ".json,image/png",
+			accept: ".json,image/png,.latent",
 			style: { display: "none" },
 			parent: document.body,
 			onchange: () => {