From de172f8be716ab8a28c202d8162e62a184f39acd Mon Sep 17 00:00:00 2001
From: Gorka Eguileor <gorka@eguileor.com>
Date: Sat, 6 Apr 2024 18:10:17 +0200
Subject: [PATCH] Improve A1111 metadata parsing (#3216)

* A1111 import: Set VAE name

This patch sets the VAE name for the `VAELoader` when present in the png
metadata.

* A1111 import: Skip all hashes

When importing from A1111 the parsing assumes that values of a key will
never contain a ":", which is not correct.

There are 2 cases where we can have ":" in the value:

- Inside a string. E.g.:
  Lora hashes: "xl_more_art-full_v1: fe3b4816be83, add-detail-xl: 9c783c8ce46c"

- When the value is a json dictionary. E.g.:
  Hashes: {"vae": "63aeecb90f", "embed:negativeXL_D": "fff5d51ab6"}

This patch changes how we parse the metadata to take those 2 cases into
account and also skips the following additional keys that are present in
some Forge images:

- Version
- VAE hash
- TI hashes
- Lora hashes
- Hashes

* A1111 import: Parse Hires steps

This patch parses the `Hires steps` parameter that is part of the High
Resolution Upscale configuration when it  is present, and fallbacks to
the one from the `samplerNode` (like the code currently does) if it's
not present.
---
 web/scripts/pnginfo.js | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/web/scripts/pnginfo.js b/web/scripts/pnginfo.js
index 1696092..7132fb6 100644
--- a/web/scripts/pnginfo.js
+++ b/web/scripts/pnginfo.js
@@ -170,9 +170,12 @@ export async function importA1111(graph, parameters) {
 		const opts = parameters
 			.substr(p)
 			.split("\n")[1]
-			.split(",")
+			.match(new RegExp("\\s*([^:]+:\\s*([^\"\\{].*?|\".*?\"|\\{.*?\\}))\\s*(,|$)", "g"))
 			.reduce((p, n) => {
 				const s = n.split(":");
+				if (s[1].endsWith(',')) {
+					s[1] = s[1].substr(0, s[1].length -1);
+				}
 				p[s[0].trim().toLowerCase()] = s[1].trim();
 				return p;
 			}, {});
@@ -191,6 +194,7 @@ export async function importA1111(graph, parameters) {
 			const vaeLoaderNode = LiteGraph.createNode("VAELoader");
 			const saveNode = LiteGraph.createNode("SaveImage");
 			let hrSamplerNode = null;
+			let hrSteps = null;
 
 			const ceil64 = (v) => Math.ceil(v / 64) * 64;
 
@@ -290,6 +294,9 @@ export async function importA1111(graph, parameters) {
 				model(v) {
 					setWidgetValue(ckptNode, "ckpt_name", v, true);
 				},
+				"vae"(v) {
+					setWidgetValue(vaeLoaderNode, "vae_name", v, true);
+				},
 				"cfg scale"(v) {
 					setWidgetValue(samplerNode, "cfg", +v);
 				},
@@ -316,6 +323,7 @@ export async function importA1111(graph, parameters) {
 					const h = ceil64(+wxh[1]);
 					const hrUp = popOpt("hires upscale");
 					const hrSz = popOpt("hires resize");
+					hrSteps = popOpt("hires steps");
 					let hrMethod = popOpt("hires upscaler");
 
 					setWidgetValue(imageNode, "width", w);
@@ -398,7 +406,7 @@ export async function importA1111(graph, parameters) {
 			}
 
 			if (hrSamplerNode) {
-				setWidgetValue(hrSamplerNode, "steps", getWidget(samplerNode, "steps").value);
+				setWidgetValue(hrSamplerNode, "steps", hrSteps? +hrSteps : getWidget(samplerNode, "steps").value);
 				setWidgetValue(hrSamplerNode, "cfg", getWidget(samplerNode, "cfg").value);
 				setWidgetValue(hrSamplerNode, "scheduler", getWidget(samplerNode, "scheduler").value);
 				setWidgetValue(hrSamplerNode, "sampler_name", getWidget(samplerNode, "sampler_name").value);
@@ -415,7 +423,7 @@ export async function importA1111(graph, parameters) {
 
 			graph.arrange();
 
-			for (const opt of ["model hash", "ensd"]) {
+			for (const opt of ["model hash", "ensd", "version", "vae hash", "ti hashes", "lora hashes", "hashes"]) {
 				delete opts[opt];
 			}