forked from pinks/eris
1
0
Fork 0
eris/bot/parsePngInfo.ts

164 lines
4.7 KiB
TypeScript

import { decode } from "png_chunk_text";
import extractChunks from "png_chunks_extract";
export function getPngInfo(pngData: Uint8Array): string | undefined {
return extractChunks(pngData)
.filter((chunk) => chunk.name === "tEXt")
.map((chunk) => decode(chunk.data))
.find((textChunk) => textChunk.keyword === "parameters")
?.text;
}
export interface PngInfo {
prompt: string;
negative_prompt: string;
steps: number;
cfg_scale: number;
width: number;
height: number;
sampler_name: string;
seed: number;
denoising_strength: number;
}
interface PngInfoExtra extends PngInfo {
upscale?: number;
file?: boolean;
}
export function parsePngInfo(
pngInfo: string,
baseParams?: Partial<PngInfo>,
shouldParseSeed?: boolean,
): Partial<PngInfo> {
const tags = pngInfo.split(/[,;]+|\.+\s|\n/u);
let part: "prompt" | "negative_prompt" | "params" = "prompt";
const params: Partial<PngInfoExtra> = {};
const prompt: string[] = [];
const negativePrompt: string[] = [];
for (const tag of tags) {
const paramValuePair = tag.trim().match(/^(\w+\s*\w*):\s+(.*)$/u);
if (paramValuePair) {
const [_match, param = "", value = ""] = paramValuePair;
switch (param.replace(/\s+/u, "").toLowerCase()) {
case "positiveprompt":
case "positive":
case "prompt":
case "pos":
part = "prompt";
prompt.push(value.trim());
break;
case "negativeprompt":
case "negative":
case "neg":
part = "negative_prompt";
negativePrompt.push(value.trim());
break;
case "steps":
case "cycles": {
part = "params";
const steps = Number(value.trim());
if (steps > 0) params.steps = Math.min(steps, 50);
break;
}
case "cfgscale":
case "cfg":
case "detail": {
part = "params";
const cfgScale = Number(value.trim());
if (cfgScale > 0) params.cfg_scale = Math.min(cfgScale, 20);
break;
}
case "size":
case "resolution": {
part = "params";
const [width = 0, height = 0] = value.trim()
.split(/\s*[x,]\s*/u, 2)
.map((v) => v.trim())
.map(Number);
if (width > 0 && height > 0) {
params.width = Math.min(width, 2048);
params.height = Math.min(height, 2048);
}
break;
}
case "upscale":
case "scale": {
part = "params";
const upscale = Number(value.trim());
if (upscale > 0) params.upscale = Math.min(upscale, 2);
break;
}
case "denoisingstrength":
case "denoising":
case "denoise": {
part = "params";
// allow percent or decimal
let denoisingStrength: number;
if (value.trim().endsWith("%")) {
denoisingStrength = Number(value.trim().slice(0, -1).trim()) / 100;
} else {
denoisingStrength = Number(value.trim());
}
denoisingStrength = Math.min(Math.max(denoisingStrength, 0), 1);
params.denoising_strength = denoisingStrength;
break;
}
case "seed": {
part = "params";
if (shouldParseSeed) {
const seed = Number(value.trim());
if (Number.isFinite(seed)) {
params.seed = seed;
}
}
break;
}
case "file": {
part = "params";
const file = value === "true" || value === "1" || value === "yes";
params.file = file;
break;
}
case "model":
case "modelhash":
case "modelname":
case "sampler":
part = "params";
// ignore for now
break;
default:
break;
}
} else if (tag.trim().length > 0) {
switch (part) {
case "prompt":
prompt.push(tag.trim());
break;
case "negative_prompt":
negativePrompt.push(tag.trim());
break;
default:
break;
}
}
}
if (prompt.length > 0) params.prompt = prompt.join(", ");
if (negativePrompt.length > 0) params.negative_prompt = negativePrompt.join(", ");
// handle upscale
if (params.upscale && baseParams?.width && baseParams?.height) {
params.width = baseParams.width * params.upscale;
params.height = baseParams.height * params.upscale;
}
return {
...baseParams,
...params,
prompt: [baseParams?.prompt, params.prompt]
.filter(Boolean).join("\n"),
negative_prompt: [baseParams?.negative_prompt, params.negative_prompt]
.filter(Boolean).join("\n"),
};
}