Import character icon from CHARX

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<Buffer>} 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<Buffer>} 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 : '';
 }
 

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);

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>} Buffer containing the extracted file
+ * @returns {Promise<Buffer|null>} 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));
     }));
 }