rixty
/
SillyTavern
mirror of https://github.com/SillyTavern/SillyTavern.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Extract PNG read/write methods

This commit is contained in:
Cohee 2024-02-16 18:03:56 +02:00
parent 0da0d490c7
commit c8b0030f6e
3 changed files with 134 additions and 43 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

89
src/character-card-parser.js
View File

@ -1,41 +1,80 @@
const fs = require('fs');
const encode = require('png-chunks-encode');
const extract = require('png-chunks-extract');
const PNGtext = require('png-chunk-text');
const parse = async (cardUrl, format) => {
/**
* Writes Character metadata to a PNG image buffer.
* @param {Buffer} image PNG image buffer
* @param {string} data Character data to write
* @returns {Buffer} PNG image buffer with metadata
*/
const write = (image, data) => {
const chunks = extract(image);
const tEXtChunks = chunks.filter(chunk => chunk.name === 'tEXt');
// Remove all existing tEXt chunks
for (let tEXtChunk of tEXtChunks) {
chunks.splice(chunks.indexOf(tEXtChunk), 1);
}
// Add new chunks before the IEND chunk
const base64EncodedData = Buffer.from(data, 'utf8').toString('base64');
chunks.splice(-1, 0, PNGtext.encode('chara', base64EncodedData));
const newBuffer = Buffer.from(encode(chunks));
return newBuffer;
};
/**
* Reads Character metadata from a PNG image buffer.
* @param {Buffer} image PNG image buffer
* @returns {string} Character data
*/
const read = (image) => {
const chunks = extract(image);
const textChunks = chunks.filter(function (chunk) {
return chunk.name === 'tEXt';
}).map(function (chunk) {
return PNGtext.decode(chunk.data);
});
if (textChunks.length === 0) {
console.error('PNG metadata does not contain any text chunks.');
throw new Error('No PNG metadata.');
}
let index = textChunks.findIndex((chunk) => chunk.keyword.toLowerCase() == 'chara');
if (index === -1) {
console.error('PNG metadata does not contain any character data.');
throw new Error('No PNG metadata.');
}
return Buffer.from(textChunks[index].text, 'base64').toString('utf8');
};
/**
* Parses a card image and returns the character metadata.
* @param {string} cardUrl Path to the card image
* @param {string} format File format
* @returns {string} Character data
*/
const parse = (cardUrl, format) => {
let fileFormat = format === undefined ? 'png' : format;
switch (fileFormat) {
case 'png': {
const buffer = fs.readFileSync(cardUrl);
const chunks = extract(buffer);
const textChunks = chunks.filter(function (chunk) {
return chunk.name === 'tEXt';
}).map(function (chunk) {
return PNGtext.decode(chunk.data);
});
if (textChunks.length === 0) {
console.error('PNG metadata does not contain any text chunks.');
throw new Error('No PNG metadata.');
}
let index = textChunks.findIndex((chunk) => chunk.keyword.toLowerCase() == 'chara');
if (index === -1) {
console.error('PNG metadata does not contain any character data.');
throw new Error('No PNG metadata.');
}
return Buffer.from(textChunks[index].text, 'base64').toString('utf8');
return read(buffer);
}
default:
break;
}
throw new Error('Unsupported format');
};
module.exports = {
parse: parse,
parse,
write,
read,
};

21
src/endpoints/characters.js
View File

@ -7,9 +7,6 @@ const writeFileAtomicSync = require('write-file-atomic').sync;
const yaml = require('yaml');
const _ = require('lodash');
const encode = require('png-chunks-encode');
const extract = require('png-chunks-extract');
const PNGtext = require('png-chunk-text');
const jimp = require('jimp');
const { DIRECTORIES, UPLOADS_PATH, AVATAR_WIDTH, AVATAR_HEIGHT } = require('../constants');
@ -33,7 +30,7 @@ const characterDataCache = new Map();
* @param {string} input_format - 'png'
* @returns {Promise<string | undefined>} - Character card data
*/
async function charaRead(img_url, input_format) {
async function charaRead(img_url, input_format = 'png') {
const stat = fs.statSync(img_url);
const cacheKey = `${img_url}-${stat.mtimeMs}`;
if (characterDataCache.has(cacheKey)) {
@ -59,22 +56,12 @@ async function charaWrite(img_url, data, target_img, response = undefined, mes =
}
}
// Read the image, resize, and save it as a PNG into the buffer
const image = await tryReadImage(img_url, crop);
const inputImage = await tryReadImage(img_url, crop);
// Get the chunks
const chunks = extract(image);
const tEXtChunks = chunks.filter(chunk => chunk.name === 'tEXt');
const outputImage = characterCardParser.write(inputImage, data);
// Remove all existing tEXt chunks
for (let tEXtChunk of tEXtChunks) {
chunks.splice(chunks.indexOf(tEXtChunk), 1);
}
// Add new chunks before the IEND chunk
const base64EncodedData = Buffer.from(data, 'utf8').toString('base64');
chunks.splice(-1, 0, PNGtext.encode('chara', base64EncodedData));
//chunks.splice(-1, 0, text.encode('lorem', 'ipsum'));
writeFileAtomicSync(DIRECTORIES.characters + target_img + '.png', Buffer.from(encode(chunks)));
writeFileAtomicSync(DIRECTORIES.characters + target_img + '.png', outputImage);
if (response !== undefined) response.send(mes);
return true;
} catch (err) {

67
src/endpoints/content-manager.js
View File

@ -10,6 +10,7 @@ const contentLogPath = path.join(contentDirectory, 'content.log');
const contentIndexPath = path.join(contentDirectory, 'index.json');
const { DIRECTORIES } = require('../constants');
const presetFolders = [DIRECTORIES.koboldAI_Settings, DIRECTORIES.openAI_Settings, DIRECTORIES.novelAI_Settings, DIRECTORIES.textGen_Settings];
const characterCardParser = require('../character-card-parser.js');
/**
* Gets the default presets from the content directory.
@ -219,6 +220,60 @@ async function downloadChubCharacter(id) {
return { buffer, fileName, fileType };
}
/**
* Downloads a character card from the Pygsite.
* @param {string} id UUID of the character
* @returns {Promise<{buffer: Buffer, fileName: string, fileType: string}>}
*/
async function downloadPygmalionCharacter(id) {
const result = await fetch('https://server.pygmalion.chat/galatea.v1.PublicCharacterService/CharacterExport', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ 'character_id': id }),
});
if (!result.ok) {
const text = await result.text();
console.log('Pygsite returned error', result.statusText, text);
throw new Error('Failed to download character');
}
const jsonData = await result.json();
const card = jsonData?.card;
if (!card || typeof card !== 'object') {
console.error('Pygsite returned invalid character data', jsonData);
throw new Error('Failed to download character');
}
try {
const avatarUrl = card?.data?.avatar;
if (!avatarUrl) {
console.error('Pygsite character does not have an avatar', card);
throw new Error('Failed to download avatar');
}
const avatarResult = await fetch(avatarUrl);
const avatarBuffer = await avatarResult.buffer();
const cardBuffer = characterCardParser.write(avatarBuffer, JSON.stringify(card));
return {
buffer: cardBuffer,
fileName: `${sanitize(id)}.png`,
fileType: 'image/png',
};
} catch(e) {
console.error('Failed to download avatar, using JSON instead', e);
return {
buffer: Buffer.from(JSON.stringify(jsonData)),
fileName: `${sanitize(id)}.json`,
fileType: 'application/json',
};
}
}
/**
*
* @param {String} str
@ -317,7 +372,17 @@ router.post('/import', jsonParser, async (request, response) => {
let type;
const isJannnyContent = url.includes('janitorai');
if (isJannnyContent) {
const isPygmalionContent = url.includes('pygmalion.chat');
if (isPygmalionContent) {
const uuid = url.split('/').pop();
if (!uuid) {
return response.sendStatus(404);
}
type = 'character';
result = await downloadPygmalionCharacter(uuid);
} else if (isJannnyContent) {
const uuid = parseJannyUrl(url);
if (!uuid) {
return response.sendStatus(404);