rixty/SillyTavern
1
0
Fork 0
mirror of https://github.com/SillyTavern/SillyTavern.git synced 2025-06-05 21:59:27 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'staging' into InfermaticAI

Browse Source
This commit is contained in:
NWilson
2024-02-22 08:32:10 -06:00
parent 7c12c836f2 beb5e470a2
commit f569424f3e
10 changed files with 349 additions and 76 deletions
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,
};

47
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) {
@@ -152,13 +139,13 @@ const processCharacter = async (item, i) => {
const img_data = await charaRead(DIRECTORIES.characters + item);
if (img_data === undefined) throw new Error('Failed to read character file');
let jsonObject = getCharaCardV2(JSON.parse(img_data));
let jsonObject = getCharaCardV2(JSON.parse(img_data), false);
jsonObject.avatar = item;
characters[i] = jsonObject;
characters[i]['json_data'] = img_data;
const charStat = fs.statSync(path.join(DIRECTORIES.characters, item));
characters[i]['date_added'] = charStat.birthtimeMs;
characters[i]['create_date'] = jsonObject['create_date'] || humanizedISO8601DateTime(charStat.birthtimeMs);
characters[i]['date_added'] = charStat.ctimeMs;
characters[i]['create_date'] = jsonObject['create_date'] || humanizedISO8601DateTime(charStat.ctimeMs);
const char_dir = path.join(DIRECTORIES.chats, item.replace('.png', ''));
const { chatSize, dateLastChat } = calculateChatSize(char_dir);
@@ -183,15 +170,30 @@ const processCharacter = async (item, i) => {
}
};
function getCharaCardV2(jsonObject) {
/**
* Convert a character object to Spec V2 format.
* @param {object} jsonObject Character object
* @param {boolean} hoistDate Will set the chat and create_date fields to the current date if they are missing
* @returns {object} Character object in Spec V2 format
*/
function getCharaCardV2(jsonObject, hoistDate = true) {
if (jsonObject.spec === undefined) {
jsonObject = convertToV2(jsonObject);
if (hoistDate && !jsonObject.create_date) {
jsonObject.create_date = humanizedISO8601DateTime();
}
} else {
jsonObject = readFromV2(jsonObject);
}
return jsonObject;
}
/**
* Convert a character object to Spec V2 format.
* @param {object} char Character object
* @returns {object} Character object in Spec V2 format
*/
function convertToV2(char) {
// Simulate incoming data from frontend form
const result = charaFormatData({
@@ -212,7 +214,8 @@ function convertToV2(char) {
});
result.chat = char.chat ?? humanizedISO8601DateTime();
result.create_date = char.create_date ?? humanizedISO8601DateTime();
result.create_date = char.create_date;
return result;
}

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,56 @@ 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/api/export/character/${id}/v2`);
if (!result.ok) {
const text = await result.text();
console.log('Pygsite returned error', result.status, text);
throw new Error('Failed to download character');
}
const jsonData = await result.json();
const characterData = jsonData?.character;
if (!characterData || typeof characterData !== 'object') {
console.error('Pygsite returned invalid character data', jsonData);
throw new Error('Failed to download character');
}
try {
const avatarUrl = characterData?.data?.avatar;
if (!avatarUrl) {
console.error('Pygsite character does not have an avatar', characterData);
throw new Error('Failed to download avatar');
}
const avatarResult = await fetch(avatarUrl);
const avatarBuffer = await avatarResult.buffer();
const cardBuffer = characterCardParser.write(avatarBuffer, JSON.stringify(characterData));
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
@@ -294,7 +345,7 @@ async function downloadJannyCharacter(uuid) {
* @param {String} url
* @returns {String | null } UUID of the character
*/
function parseJannyUrl(url) {
function getUuidFromUrl(url) {
// Extract UUID from URL
const uuidRegex = /[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}/;
const matches = url.match(uuidRegex);
@@ -317,8 +368,18 @@ router.post('/import', jsonParser, async (request, response) => {
let type;
const isJannnyContent = url.includes('janitorai');
if (isJannnyContent) {
const uuid = parseJannyUrl(url);
const isPygmalionContent = url.includes('pygmalion.chat');
if (isPygmalionContent) {
const uuid = getUuidFromUrl(url);
if (!uuid) {
return response.sendStatus(404);
}
type = 'character';
result = await downloadPygmalionCharacter(uuid);
} else if (isJannnyContent) {
const uuid = getUuidFromUrl(url);
if (!uuid) {
return response.sendStatus(404);
}

61
src/util.js
View File

@@ -365,7 +365,7 @@ function getImages(path) {
/**
* Pipe a fetch() response to an Express.js Response, including status code.
* @param {import('node-fetch').Response} from The Fetch API response to pipe from.
* @param {Express.Response} to The Express response to pipe to.
* @param {import('express').Response} to The Express response to pipe to.
*/
function forwardFetchResponse(from, to) {
let statusCode = from.status;
@@ -399,6 +399,64 @@ function forwardFetchResponse(from, to) {
});
}
/**
* Makes an HTTP/2 request to the specified endpoint.
*
* @deprecated Use `node-fetch` if possible.
* @param {string} endpoint URL to make the request to
* @param {string} method HTTP method to use
* @param {string} body Request body
* @param {object} headers Request headers
* @returns {Promise<string>} Response body
*/
function makeHttp2Request(endpoint, method, body, headers) {
return new Promise((resolve, reject) => {
try {
const http2 = require('http2');
const url = new URL(endpoint);
const client = http2.connect(url.origin);
const req = client.request({
':method': method,
':path': url.pathname,
...headers,
});
req.setEncoding('utf8');
req.on('response', (headers) => {
const status = Number(headers[':status']);
if (status < 200 || status >= 300) {
reject(new Error(`Request failed with status ${status}`));
}
let data = '';
req.on('data', (chunk) => {
data += chunk;
});
req.on('end', () => {
console.log(data);
resolve(data);
});
});
req.on('error', (err) => {
reject(err);
});
if (body) {
req.write(body);
}
req.end();
} catch (e) {
reject(e);
}
});
}
/**
* Adds YAML-serialized object to the object.
* @param {object} obj Object
@@ -547,4 +605,5 @@ module.exports = {
excludeKeysByYaml,
trimV1,
Cache,
makeHttp2Request,
};