From 60b7164c283ef12bc0afc5d37238f8ef3beb6307 Mon Sep 17 00:00:00 2001 From: Cohee <18619528+Cohee1207@users.noreply.github.com> Date: Sat, 8 Jun 2024 22:37:17 +0300 Subject: [PATCH] Import character icon from CHARX --- src/endpoints/characters.js | 60 ++++++++++++++++++++++++++++++++++--- src/endpoints/novelai.js | 11 +++++++ src/util.js | 3 +- 3 files changed, 69 insertions(+), 5 deletions(-) diff --git a/src/endpoints/characters.js b/src/endpoints/characters.js index 8e663db56..c5c15b2e3 100644 --- a/src/endpoints/characters.js +++ b/src/endpoints/characters.js @@ -44,7 +44,7 @@ async function readCharacterData(inputFile, inputFormat = 'png') { /** * Writes the character card to the specified image file. - * @param {string} inputFile - Path to the image file + * @param {string|Buffer} inputFile - Path to the image file or image buffer * @param {string} data - Character card data * @param {string} outputFile - Target image file name * @param {import('express').Request} request - Express request obejct @@ -60,8 +60,20 @@ async function writeCharacterData(inputFile, data, outputFile, request, crop = u break; } } - // Read the image, resize, and save it as a PNG into the buffer - const inputImage = await tryReadImage(inputFile, crop); + + /** + * Read the image, resize, and save it as a PNG into the buffer. + * @returns {Promise} Image buffer + */ + function getInputImage() { + if (Buffer.isBuffer(inputFile)) { + return parseImageBuffer(inputFile, crop); + } + + return tryReadImage(inputFile, crop); + } + + const inputImage = await getInputImage(); // Get the chunks const outputImage = characterCardParser.write(inputImage, data); @@ -84,6 +96,32 @@ async function writeCharacterData(inputFile, data, outputFile, request, crop = u * @property {boolean} want_resize Resize the image to the standard avatar size */ +/** + * Parses an image buffer and applies crop if defined. + * @param {Buffer} buffer Buffer of the image + * @param {Crop|undefined} [crop] Crop parameters + * @returns {Promise} Image buffer + */ +async function parseImageBuffer(buffer, crop) { + const image = await jimp.read(buffer); + let finalWidth = image.bitmap.width, finalHeight = image.bitmap.height; + + // Apply crop if defined + if (typeof crop == 'object' && [crop.x, crop.y, crop.width, crop.height].every(x => typeof x === 'number')) { + image.crop(crop.x, crop.y, crop.width, crop.height); + // Apply standard resize if requested + if (crop.want_resize) { + finalWidth = AVATAR_WIDTH; + finalHeight = AVATAR_HEIGHT; + } else { + finalWidth = crop.width; + finalHeight = crop.height; + } + } + + return image.cover(finalWidth, finalHeight).getBufferAsync(jimp.MIME_PNG); +} + /** * Reads an image file and applies crop if defined. * @param {string} imgPath Path to the image file @@ -509,11 +547,25 @@ async function importFromCharX(uploadPath, { request }) { throw new Error('Invalid CharX card file: missing spec field'); } + /** @type {string|Buffer} */ + let avatar = defaultAvatarPath; + const assets = _.get(card, 'data.assets'); + if (Array.isArray(assets) && assets.length) { + for (const asset of assets.filter(x => x.type === 'icon' && typeof x.uri === 'string')) { + const pathNoProtocol = String(asset.uri.replace(/^(?:\/\/|[^/]+)*\//, '')); + const buffer = await extractFileFromZipBuffer(data, pathNoProtocol); + if (buffer) { + avatar = buffer; + break; + } + } + } + unsetFavFlag(card); card['create_date'] = humanizedISO8601DateTime(); card.name = sanitize(card.name); const fileName = getPngName(card.name, request.user.directories); - const result = await writeCharacterData(defaultAvatarPath, JSON.stringify(card), fileName, request); + const result = await writeCharacterData(avatar, JSON.stringify(card), fileName, request); return result ? fileName : ''; } diff --git a/src/endpoints/novelai.js b/src/endpoints/novelai.js index e26bfe34a..abbfe3ef6 100644 --- a/src/endpoints/novelai.js +++ b/src/endpoints/novelai.js @@ -281,6 +281,12 @@ router.post('/generate-image', jsonParser, async (request, response) => { const archiveBuffer = await generateResult.arrayBuffer(); const imageBuffer = await extractFileFromZipBuffer(archiveBuffer, '.png'); + + if (!imageBuffer) { + console.warn('NovelAI generated an image, but the PNG file was not found.'); + return response.sendStatus(500); + } + const originalBase64 = imageBuffer.toString('base64'); // No upscaling @@ -311,6 +317,11 @@ router.post('/generate-image', jsonParser, async (request, response) => { const upscaledArchiveBuffer = await upscaleResult.arrayBuffer(); const upscaledImageBuffer = await extractFileFromZipBuffer(upscaledArchiveBuffer, '.png'); + + if (!upscaledImageBuffer) { + throw new Error('NovelAI upscaled an image, but the PNG file was not found.'); + } + const upscaledBase64 = upscaledImageBuffer.toString('base64'); return response.send(upscaledBase64); diff --git a/src/util.js b/src/util.js index 8ccc7851c..326ad6a36 100644 --- a/src/util.js +++ b/src/util.js @@ -139,7 +139,7 @@ function getHexString(length) { * Extracts a file with given extension from an ArrayBuffer containing a ZIP archive. * @param {ArrayBuffer} archiveBuffer Buffer containing a ZIP archive * @param {string} fileExtension File extension to look for - * @returns {Promise} Buffer containing the extracted file + * @returns {Promise} Buffer containing the extracted file. Null if the file was not found. */ async function extractFileFromZipBuffer(archiveBuffer, fileExtension) { return await new Promise((resolve, reject) => yauzl.fromBuffer(Buffer.from(archiveBuffer), { lazyEntries: true }, (err, zipfile) => { @@ -171,6 +171,7 @@ async function extractFileFromZipBuffer(archiveBuffer, fileExtension) { zipfile.readEntry(); } }); + zipfile.on('end', () => resolve(null)); })); }