Files
ComfyUI_frontend/src/scripts/pnginfo.ts
Chenlei Hu acdaa6a594 Format all code / Add pre-commit format hook (#81)
* Add format-guard

* Format code
2024-07-02 13:22:37 -04:00

588 lines
18 KiB
TypeScript

import { api } from "./api";
export function getPngMetadata(file) {
return new Promise<Record<string, string>>((r) => {
const reader = new FileReader();
reader.onload = (event) => {
// Get the PNG data as a Uint8Array
const pngData = new Uint8Array(event.target.result as ArrayBuffer);
const dataView = new DataView(pngData.buffer);
// Check that the PNG signature is present
if (dataView.getUint32(0) !== 0x89504e47) {
console.error("Not a valid PNG file");
r({});
return;
}
// Start searching for chunks after the PNG signature
let offset = 8;
let txt_chunks: Record<string, string> = {};
// Loop through the chunks in the PNG file
while (offset < pngData.length) {
// Get the length of the chunk
const length = dataView.getUint32(offset);
// Get the chunk type
const type = String.fromCharCode(
...pngData.slice(offset + 4, offset + 8)
);
if (type === "tEXt" || type == "comf" || type === "iTXt") {
// Get the keyword
let keyword_end = offset + 8;
while (pngData[keyword_end] !== 0) {
keyword_end++;
}
const keyword = String.fromCharCode(
...pngData.slice(offset + 8, keyword_end)
);
// Get the text
const contentArraySegment = pngData.slice(
keyword_end + 1,
offset + 8 + length
);
const contentJson = new TextDecoder("utf-8").decode(
contentArraySegment
);
txt_chunks[keyword] = contentJson;
}
offset += 12 + length;
}
r(txt_chunks);
};
reader.readAsArrayBuffer(file);
});
}
function parseExifData(exifData) {
// Check for the correct TIFF header (0x4949 for little-endian or 0x4D4D for big-endian)
const isLittleEndian = new Uint16Array(exifData.slice(0, 2))[0] === 0x4949;
// Function to read 16-bit and 32-bit integers from binary data
function readInt(offset, isLittleEndian, length) {
let arr = exifData.slice(offset, offset + length);
if (length === 2) {
return new DataView(arr.buffer, arr.byteOffset, arr.byteLength).getUint16(
0,
isLittleEndian
);
} else if (length === 4) {
return new DataView(arr.buffer, arr.byteOffset, arr.byteLength).getUint32(
0,
isLittleEndian
);
}
}
// Read the offset to the first IFD (Image File Directory)
const ifdOffset = readInt(4, isLittleEndian, 4);
function parseIFD(offset) {
const numEntries = readInt(offset, isLittleEndian, 2);
const result = {};
for (let i = 0; i < numEntries; i++) {
const entryOffset = offset + 2 + i * 12;
const tag = readInt(entryOffset, isLittleEndian, 2);
const type = readInt(entryOffset + 2, isLittleEndian, 2);
const numValues = readInt(entryOffset + 4, isLittleEndian, 4);
const valueOffset = readInt(entryOffset + 8, isLittleEndian, 4);
// Read the value(s) based on the data type
let value;
if (type === 2) {
// ASCII string
value = String.fromCharCode(
...exifData.slice(valueOffset, valueOffset + numValues - 1)
);
}
result[tag] = value;
}
return result;
}
// Parse the first IFD
const ifdData = parseIFD(ifdOffset);
return ifdData;
}
function splitValues(input) {
var output = {};
for (var key in input) {
var value = input[key];
var splitValues = value.split(":", 2);
output[splitValues[0]] = splitValues[1];
}
return output;
}
export function getWebpMetadata(file) {
return new Promise<Record<string, string>>((r) => {
const reader = new FileReader();
reader.onload = (event) => {
const webp = new Uint8Array(event.target.result as ArrayBuffer);
const dataView = new DataView(webp.buffer);
// Check that the WEBP signature is present
if (
dataView.getUint32(0) !== 0x52494646 ||
dataView.getUint32(8) !== 0x57454250
) {
console.error("Not a valid WEBP file");
r({});
return;
}
// Start searching for chunks after the WEBP signature
let offset = 12;
let txt_chunks = {};
// Loop through the chunks in the WEBP file
while (offset < webp.length) {
const chunk_length = dataView.getUint32(offset + 4, true);
const chunk_type = String.fromCharCode(
...webp.slice(offset, offset + 4)
);
if (chunk_type === "EXIF") {
if (
String.fromCharCode(...webp.slice(offset + 8, offset + 8 + 6)) ==
"Exif\0\0"
) {
offset += 6;
}
let data = parseExifData(
webp.slice(offset + 8, offset + 8 + chunk_length)
);
for (var key in data) {
var value = data[key] as string;
let index = value.indexOf(":");
txt_chunks[value.slice(0, index)] = value.slice(index + 1);
}
}
offset += 8 + chunk_length;
}
r(txt_chunks);
};
reader.readAsArrayBuffer(file);
});
}
export function getLatentMetadata(file) {
return new Promise((r) => {
const reader = new FileReader();
reader.onload = (event) => {
const safetensorsData = new Uint8Array(
event.target.result as ArrayBuffer
);
const dataView = new DataView(safetensorsData.buffer);
let header_size = dataView.getUint32(0, true);
let offset = 8;
let header = JSON.parse(
new TextDecoder().decode(
safetensorsData.slice(offset, offset + header_size)
)
);
r(header.__metadata__);
};
var slice = file.slice(0, 1024 * 1024 * 4);
reader.readAsArrayBuffer(slice);
});
}
function getString(dataView: DataView, offset: number, length: number): string {
let string = "";
for (let i = 0; i < length; i++) {
string += String.fromCharCode(dataView.getUint8(offset + i));
}
return string;
}
// Function to parse the Vorbis Comment block
function parseVorbisComment(dataView: DataView): Record<string, string> {
let offset = 0;
const vendorLength = dataView.getUint32(offset, true);
offset += 4;
const vendorString = getString(dataView, offset, vendorLength);
offset += vendorLength;
const userCommentListLength = dataView.getUint32(offset, true);
offset += 4;
const comments = {};
for (let i = 0; i < userCommentListLength; i++) {
const commentLength = dataView.getUint32(offset, true);
offset += 4;
const comment = getString(dataView, offset, commentLength);
offset += commentLength;
const [key, value] = comment.split("=");
comments[key] = value;
}
return comments;
}
// Function to read a FLAC file and parse Vorbis comments
export function getFlacMetadata(file: Blob): Promise<Record<string, string>> {
return new Promise((r) => {
const reader = new FileReader();
reader.onload = function (event) {
const arrayBuffer = event.target.result as ArrayBuffer;
const dataView = new DataView(arrayBuffer);
// Verify the FLAC signature
const signature = String.fromCharCode(
...new Uint8Array(arrayBuffer, 0, 4)
);
if (signature !== "fLaC") {
console.error("Not a valid FLAC file");
return;
}
// Parse metadata blocks
let offset = 4;
let vorbisComment = null;
while (offset < dataView.byteLength) {
const isLastBlock = dataView.getUint8(offset) & 0x80;
const blockType = dataView.getUint8(offset) & 0x7f;
const blockSize = dataView.getUint32(offset, false) & 0xffffff;
offset += 4;
if (blockType === 4) {
// Vorbis Comment block type
vorbisComment = parseVorbisComment(
new DataView(arrayBuffer, offset, blockSize)
);
}
offset += blockSize;
if (isLastBlock) break;
}
r(vorbisComment);
};
reader.readAsArrayBuffer(file);
});
}
export async function importA1111(graph, parameters) {
const p = parameters.lastIndexOf("\nSteps:");
if (p > -1) {
const embeddings = await api.getEmbeddings();
const opts = parameters
.substr(p)
.split("\n")[1]
.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;
}, {});
const p2 = parameters.lastIndexOf("\nNegative prompt:", p);
if (p2 > -1) {
let positive = parameters.substr(0, p2).trim();
let negative = parameters.substring(p2 + 18, p).trim();
const ckptNode = LiteGraph.createNode("CheckpointLoaderSimple");
const clipSkipNode = LiteGraph.createNode("CLIPSetLastLayer");
const positiveNode = LiteGraph.createNode("CLIPTextEncode");
const negativeNode = LiteGraph.createNode("CLIPTextEncode");
const samplerNode = LiteGraph.createNode("KSampler");
const imageNode = LiteGraph.createNode("EmptyLatentImage");
const vaeNode = LiteGraph.createNode("VAEDecode");
const vaeLoaderNode = LiteGraph.createNode("VAELoader");
const saveNode = LiteGraph.createNode("SaveImage");
let hrSamplerNode = null;
let hrSteps = null;
const ceil64 = (v) => Math.ceil(v / 64) * 64;
const getWidget = (node, name) => {
return node.widgets.find((w) => w.name === name);
};
const setWidgetValue = (node, name, value, isOptionPrefix?) => {
const w = getWidget(node, name);
if (isOptionPrefix) {
const o = w.options.values.find((w) => w.startsWith(value));
if (o) {
w.value = o;
} else {
console.warn(`Unknown value '${value}' for widget '${name}'`, node);
w.value = value;
}
} else {
w.value = value;
}
};
const createLoraNodes = (clipNode, text, prevClip, prevModel) => {
const loras = [];
text = text.replace(/<lora:([^:]+:[^>]+)>/g, function (m, c) {
const s = c.split(":");
const weight = parseFloat(s[1]);
if (isNaN(weight)) {
console.warn("Invalid LORA", m);
} else {
loras.push({ name: s[0], weight });
}
return "";
});
for (const l of loras) {
const loraNode = LiteGraph.createNode("LoraLoader");
graph.add(loraNode);
setWidgetValue(loraNode, "lora_name", l.name, true);
setWidgetValue(loraNode, "strength_model", l.weight);
setWidgetValue(loraNode, "strength_clip", l.weight);
prevModel.node.connect(prevModel.index, loraNode, 0);
prevClip.node.connect(prevClip.index, loraNode, 1);
prevModel = { node: loraNode, index: 0 };
prevClip = { node: loraNode, index: 1 };
}
prevClip.node.connect(1, clipNode, 0);
prevModel.node.connect(0, samplerNode, 0);
if (hrSamplerNode) {
prevModel.node.connect(0, hrSamplerNode, 0);
}
return { text, prevModel, prevClip };
};
const replaceEmbeddings = (text) => {
if (!embeddings.length) return text;
return text.replaceAll(
new RegExp(
"\\b(" +
embeddings
.map((e) => e.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"))
.join("\\b|\\b") +
")\\b",
"ig"
),
"embedding:$1"
);
};
const popOpt = (name) => {
const v = opts[name];
delete opts[name];
return v;
};
graph.clear();
graph.add(ckptNode);
graph.add(clipSkipNode);
graph.add(positiveNode);
graph.add(negativeNode);
graph.add(samplerNode);
graph.add(imageNode);
graph.add(vaeNode);
graph.add(vaeLoaderNode);
graph.add(saveNode);
ckptNode.connect(1, clipSkipNode, 0);
clipSkipNode.connect(0, positiveNode, 0);
clipSkipNode.connect(0, negativeNode, 0);
ckptNode.connect(0, samplerNode, 0);
positiveNode.connect(0, samplerNode, 1);
negativeNode.connect(0, samplerNode, 2);
imageNode.connect(0, samplerNode, 3);
vaeNode.connect(0, saveNode, 0);
samplerNode.connect(0, vaeNode, 0);
vaeLoaderNode.connect(0, vaeNode, 1);
const handlers = {
model(v) {
setWidgetValue(ckptNode, "ckpt_name", v, true);
},
vae(v) {
setWidgetValue(vaeLoaderNode, "vae_name", v, true);
},
"cfg scale"(v) {
setWidgetValue(samplerNode, "cfg", +v);
},
"clip skip"(v) {
setWidgetValue(clipSkipNode, "stop_at_clip_layer", -v);
},
sampler(v) {
let name = v.toLowerCase().replace("++", "pp").replaceAll(" ", "_");
if (name.includes("karras")) {
name = name.replace("karras", "").replace(/_+$/, "");
setWidgetValue(samplerNode, "scheduler", "karras");
} else {
setWidgetValue(samplerNode, "scheduler", "normal");
}
const w = getWidget(samplerNode, "sampler_name");
const o = w.options.values.find(
(w) => w === name || w === "sample_" + name
);
if (o) {
setWidgetValue(samplerNode, "sampler_name", o);
}
},
size(v) {
const wxh = v.split("x");
const w = ceil64(+wxh[0]);
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);
setWidgetValue(imageNode, "height", h);
if (hrUp || hrSz) {
let uw, uh;
if (hrUp) {
uw = w * hrUp;
uh = h * hrUp;
} else {
const s = hrSz.split("x");
uw = +s[0];
uh = +s[1];
}
let upscaleNode;
let latentNode;
if (hrMethod.startsWith("Latent")) {
latentNode = upscaleNode = LiteGraph.createNode("LatentUpscale");
graph.add(upscaleNode);
samplerNode.connect(0, upscaleNode, 0);
switch (hrMethod) {
case "Latent (nearest-exact)":
hrMethod = "nearest-exact";
break;
}
setWidgetValue(upscaleNode, "upscale_method", hrMethod, true);
} else {
const decode = LiteGraph.createNode("VAEDecodeTiled");
graph.add(decode);
samplerNode.connect(0, decode, 0);
vaeLoaderNode.connect(0, decode, 1);
const upscaleLoaderNode =
LiteGraph.createNode("UpscaleModelLoader");
graph.add(upscaleLoaderNode);
setWidgetValue(upscaleLoaderNode, "model_name", hrMethod, true);
const modelUpscaleNode = LiteGraph.createNode(
"ImageUpscaleWithModel"
);
graph.add(modelUpscaleNode);
decode.connect(0, modelUpscaleNode, 1);
upscaleLoaderNode.connect(0, modelUpscaleNode, 0);
upscaleNode = LiteGraph.createNode("ImageScale");
graph.add(upscaleNode);
modelUpscaleNode.connect(0, upscaleNode, 0);
const vaeEncodeNode = (latentNode =
LiteGraph.createNode("VAEEncodeTiled"));
graph.add(vaeEncodeNode);
upscaleNode.connect(0, vaeEncodeNode, 0);
vaeLoaderNode.connect(0, vaeEncodeNode, 1);
}
setWidgetValue(upscaleNode, "width", ceil64(uw));
setWidgetValue(upscaleNode, "height", ceil64(uh));
hrSamplerNode = LiteGraph.createNode("KSampler");
graph.add(hrSamplerNode);
ckptNode.connect(0, hrSamplerNode, 0);
positiveNode.connect(0, hrSamplerNode, 1);
negativeNode.connect(0, hrSamplerNode, 2);
latentNode.connect(0, hrSamplerNode, 3);
hrSamplerNode.connect(0, vaeNode, 0);
}
},
steps(v) {
setWidgetValue(samplerNode, "steps", +v);
},
seed(v) {
setWidgetValue(samplerNode, "seed", +v);
},
};
for (const opt in opts) {
if (opt in handlers) {
handlers[opt](popOpt(opt));
}
}
if (hrSamplerNode) {
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
);
setWidgetValue(
hrSamplerNode,
"denoise",
+(popOpt("denoising strength") || "1")
);
}
let n = createLoraNodes(
positiveNode,
positive,
{ node: clipSkipNode, index: 0 },
{ node: ckptNode, index: 0 }
);
positive = n.text;
n = createLoraNodes(negativeNode, negative, n.prevClip, n.prevModel);
negative = n.text;
setWidgetValue(positiveNode, "text", replaceEmbeddings(positive));
setWidgetValue(negativeNode, "text", replaceEmbeddings(negative));
graph.arrange();
for (const opt of [
"model hash",
"ensd",
"version",
"vae hash",
"ti hashes",
"lora hashes",
"hashes",
]) {
delete opts[opt];
}
console.warn("Unhandled parameters:", opts);
}
}
}