Skip to content

Commit

Permalink
refactor: code splitting and update support \x00\x00\x00
Browse files Browse the repository at this point in the history
  • Loading branch information
bluelovers committed May 10, 2024
1 parent f1cbee8 commit b144fd3
Show file tree
Hide file tree
Showing 17 changed files with 803 additions and 194 deletions.
11 changes: 10 additions & 1 deletion package.json
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,12 @@
"stable",
"diffusion",
"webui",
"sdwebui",
"sd",
"prompt",
"prompts",
"Stable Diffusion web UI",
"Stable Diffusion",
"create-by-yarn-tool",
"create-by-tsdx"
],
Expand Down Expand Up @@ -93,10 +99,13 @@
"ncu": "yarn-tool ncu -u",
"tsc:showConfig": "ynpx get-current-tsconfig -p"
},
"dependencies": {
"split-smartly2": "^2.0.1"
},
"devDependencies": {
"@bluelovers/tsconfig": "^1.0.35",
"@types/jest": "^29.5.12",
"@types/node": "^20.12.7"
"@types/node": "^20.12.11"
},
"packageManager": "yarn@1.22.19"
}
23 changes: 23 additions & 0 deletions src/handler.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
import { IOptionsInfoparser } from './index';

export function keyToSnakeStyle1(key: string)
{
return key.toLowerCase().replace(/ /g, '_')
}

export function handleInfoEntries(entries: readonly (readonly [string, string])[], opts?: IOptionsInfoparser)
{
const cast_to_snake = opts?.cast_to_snake;
const re = /^0\d/;

return entries.map(([key, value]) =>
{
const asNum = parseFloat(value);
const isNotNum = (re.test(value)) || isNaN(asNum) || ((value as any as number) - asNum) !== 0;

if (cast_to_snake) key = keyToSnakeStyle1(key)

const out = [key, isNotNum ? value : asNum] as const;
return out
})
}
167 changes: 66 additions & 101 deletions src/index.ts
Original file line number Diff line number Diff line change
@@ -1,117 +1,83 @@
/**
* get int32 from png compensating for endianness
*/
export function i32(a: Uint8Array,
i: number,
)
import { inputToBytes } from './utils';
import { _parseInfoLine, extractPromptAndInfoFromRaw } from './parser';
import { extractRawFromBytes } from './png';
import { handleInfoEntries } from './handler';

export interface IOptionsInfoparser
{
return new Uint32Array(new Uint8Array([...a.slice(i, i + 4)].reverse()).buffer)[0];
cast_to_snake?: boolean;
isIncludePrompts?: boolean;
}

export function infoparser(line: string)
export function infoparser(line: string, opts?: IOptionsInfoparser)
{
let output: Record<string, string> = {}
let buffer = ''
let key: string
let quotes = false
//actually no idea why there are trailing : sometimes?
if (line.at(-1) === ':') line = line.slice(0, -1)
let base = [] as ReturnType<typeof handleInfoEntries>
if (opts?.isIncludePrompts)
{
const {
prompt,
negative_prompt,
infoline,
} = extractPromptAndInfoFromRaw(line);

let c_pre: string;
base.push(['prompt', prompt]);
base.push(['negative_prompt', negative_prompt]);
line = infoline;
}

for (let i = 0; i < line.length; i++)
{
let c = line[i];
let sp = true;
return Object.fromEntries(base.concat(handleInfoEntries(_parseInfoLine(line), opts)))
}

if (c === '"')
{
quotes = !quotes
}
/**
* @example
* import fs from 'fs/promises'
* import PNGINFO from 'auto1111-pnginfo'
*
* const file = await fs.readFile('generate_waifu.png')
* const info = PNGINFO(file)
*
* console.log(info)
*/
export function PNGINFO(png: Uint8Array | string, cast_to_snake = false)
{
let bytes = inputToBytes(png) as Uint8Array;

if (!quotes)
{
if (c === ':')
{
key = buffer.trim();
buffer = ''
sp = false;
}
else if (c === ',')
{
if (key === 'Cutoff targets' && c_pre !== ']')
{
const raw = extractRawFromBytes(bytes);
if (!raw) return;

}
else
{
output[key] = buffer.slice(1);
buffer = '';
sp = false;
}
}
}
const {
raw_info,
width,
height
} = raw;

if (sp)
{
buffer += c;
}
const {
prompt,
negative_prompt,
infoline,
infoline_extra,
//lines_raw,
} = extractPromptAndInfoFromRaw(raw_info as any)

c_pre = c;
}
if (key) output[key] = buffer.slice(1)
return output
}
let data = infoparser(infoline, {
cast_to_snake
});

export function PNGINFO(png: Uint8Array | string, cast_to_snake = false)
{
let bin_str: string, bytes: Uint8Array
if (typeof Buffer !== 'undefined' && Buffer.isBuffer(png))
{
bytes = Uint8Array.from(png)
bin_str = png.toString()
let output = {
metadata: {
width,
height,
extra: infoline_extra,
raw_info
},
pnginfo: {
prompt,
negative_prompt,
...data,
},
}
else if (png instanceof Uint8Array)
{
bytes = png
bin_str = png.toString()
}
else
{
bin_str = atob(png.slice(0, 8192))
bytes = Uint8Array.from(bin_str, c => c.charCodeAt(0))
}
// @ts-ignore
const pngmagic = bytes.slice(0, 8) == '137,80,78,71,13,10,26,10'
if (!pngmagic) return
let [ihdrSize, width, height] = [i32(bytes, 8), i32(bytes, 16), i32(bytes, 20)]
let txtOffset = 8 + ihdrSize + 12
if (bin_str.slice(txtOffset + 4, txtOffset + 8) != 'tEXt') return
let txtSize = i32(bytes, txtOffset)
let raw_info = bin_str.slice(txtOffset + 8 + "parameters\u0000".length, txtOffset + 8 + txtSize)
let infolines = raw_info.split('\n')
let negative_prompt_index = infolines.findIndex(a => a.match(/^Negative prompt:/))
let prompt = infolines.splice(0, negative_prompt_index).join('\n').trim()
let steps_index = infolines.findIndex(a => a.match(/^Steps:/))
let negative_prompt = infolines.splice(0, steps_index).join('\n').slice('Negative prompt:'.length).trim()
let infoline = infolines.splice(0, 1)[0]
let data = infoparser(infoline)

data = Object.fromEntries(
Object.entries(data).map(([key, value]) =>
{
let asNum = parseFloat(value)
// @ts-ignore
let isNotNum = value - asNum
// @ts-ignore
if (cast_to_snake) key = key.toLowerCase().replaceAll(' ', '_')
let out = [key, isNotNum || isNaN(isNotNum) ? value : asNum]
return out
}),
)

let output = { width, height, prompt, negative_prompt, extra: infolines, raw_info }
return Object.assign(output, data)
return output
}

// @ts-ignore
Expand All @@ -123,7 +89,6 @@ if (process.env.TSDX_FORMAT !== 'esm')
Object.defineProperty(PNGINFO, 'default', { value: PNGINFO });

Object.defineProperty(PNGINFO, 'infoparser', { value: infoparser });
Object.defineProperty(PNGINFO, 'i32', { value: i32 });
}

export default PNGINFO
137 changes: 137 additions & 0 deletions src/parser.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,137 @@
import { splitSmartly } from 'split-smartly2';
import { _isRawVersionPlus, _splitRawToLines } from './split';

/**
* `${key}: ${value}`
*/
export function _parseLine(line: string)
{
const [, key, value] = line.match(/^([^:]+)\s*:\s*(.*)$/);
return [key, value] as const
}

/**
* Parses an info line into key-value pairs.
*
* @param infoline - The info line to parse.
* @returns An array of tuples, where each tuple contains a key-value pair.
*
* @remarks
* This function uses the `splitSmartly` function from the 'split-smartly2' package to split the info line into key-value pairs.
* The info line is expected to be in the format `${key}: ${value}`, separated by commas.
* The `splitSmartly` function is configured to handle nested brackets and trim separators.
*
* @example
* ```typescript
* const infoLine = 'key1: value1, key2: value2, key3: value3';
* const result = _parseInfoLine(infoLine);
* // result: [['key1', 'value1'], ['key2', 'value2'], ['key3', 'value3']]
* ```
*/
export function _parseInfoLine(infoline: string)
{
const entries = splitSmartly(infoline, [','], {
brackets: true,
trimSeparators: true,
}) as string[];

return entries.map(_parseLine)
}

/**
* Extracts prompt, negative prompt, info line, and extra info from a raw info string.
*
* @param raw_info - The raw info string to extract data from.
* @returns An object containing the extracted prompt, negative prompt, info line, extra info, and the original lines.
*
* @throws Will throw a TypeError if the raw info string is in Plus version and contains more than 3 lines.
*
* @remarks
* This function first checks if the raw info string is in Plus version using the `_isRawVersionPlus` function.
* It then splits the raw info string into lines using the `_splitRawToLines` function.
* Depending on the version, it extracts the prompt, negative prompt, info line, and extra info.
* If the raw info string is in Plus version, it follows a specific order to extract the data.
* If the raw info string is not in Plus version, it uses the `findIndex` method to find the indices of the negative prompt and steps,
* and then extracts the data accordingly.
* Finally, it returns an object containing the extracted data and the original lines.
*
* @example
* ```typescript
* const rawInfo = 'Prompt:...\nNegative prompt:...\nSteps:...';
* const result = extractPromptAndInfoFromRaw(rawInfo);
* // result: {
* // prompt: '...',
* // negative_prompt: '...',
* // infoline: 'Steps:...',
* // infoline_extra: [],
* // lines_raw: ['Prompt:...', 'Negative prompt:...', 'Steps:...'],
* // }
* ```
*/
export function extractPromptAndInfoFromRaw(raw_info: string)
{
const isPlus = _isRawVersionPlus(raw_info);
let lines = _splitRawToLines(raw_info);

let prompt: string = '';
let negative_prompt: string = '';
let infoline: string = '';
let infoline_extra: string[] = [];

const lines_raw = lines.slice();

if (isPlus)
{
if (lines.length > 3)
{
throw new TypeError()
}

let line = lines.pop();

if (line.startsWith('Steps: '))
{
infoline = line;
line = void 0
}

line ??= lines.pop();

if (line.startsWith('Negative prompt: '))
{
negative_prompt = line.slice('Negative prompt: '.length);
line = void 0
}

line ??= lines.pop();

prompt = line;

if (lines.length)
{
throw new TypeError()
}
}
else
{
let negative_prompt_index = lines.findIndex(a => a.match(/^Negative prompt:/))
prompt = lines.splice(0, negative_prompt_index).join('\n').trim()
let steps_index = lines.findIndex(a => a.match(/^Steps:/))
negative_prompt = lines.splice(0, steps_index).join('\n').slice('Negative prompt:'.length).trim()

infoline = lines.splice(0, 1)[0];

infoline_extra = lines;
}

prompt = prompt.replace(/\x00\x00\x00/g, '');
negative_prompt = negative_prompt.replace(/\x00\x00\x00/g, '');

return {
prompt,
negative_prompt,
infoline,
infoline_extra,
lines_raw,
}
}

0 comments on commit b144fd3

Please sign in to comment.