Spaces:
Paused
Paused
| const fs = require('fs'); | |
| const path = require('path'); | |
| const express = require('express'); | |
| const mime = require('mime-types'); | |
| const sanitize = require('sanitize-filename'); | |
| const writeFileAtomicSync = require('write-file-atomic').sync; | |
| const { getImageBuffers } = require('../util'); | |
| const { jsonParser, urlencodedParser } = require('../express-common'); | |
| /** | |
| * Gets the path to the sprites folder for the provided character name | |
| * @param {import('../users').UserDirectoryList} directories - User directories | |
| * @param {string} name - The name of the character | |
| * @param {boolean} isSubfolder - Whether the name contains a subfolder | |
| * @returns {string | null} The path to the sprites folder. Null if the name is invalid. | |
| */ | |
| function getSpritesPath(directories, name, isSubfolder) { | |
| if (isSubfolder) { | |
| const nameParts = name.split('/'); | |
| const characterName = sanitize(nameParts[0]); | |
| const subfolderName = sanitize(nameParts[1]); | |
| if (!characterName || !subfolderName) { | |
| return null; | |
| } | |
| return path.join(directories.characters, characterName, subfolderName); | |
| } | |
| name = sanitize(name); | |
| if (!name) { | |
| return null; | |
| } | |
| return path.join(directories.characters, name); | |
| } | |
| /** | |
| * Imports base64 encoded sprites from RisuAI character data. | |
| * The sprites are saved in the character's sprites folder. | |
| * The additionalAssets and emotions are removed from the data. | |
| * @param {import('../users').UserDirectoryList} directories User directories | |
| * @param {object} data RisuAI character data | |
| * @returns {void} | |
| */ | |
| function importRisuSprites(directories, data) { | |
| try { | |
| const name = data?.data?.name; | |
| const risuData = data?.data?.extensions?.risuai; | |
| // Not a Risu AI character | |
| if (!risuData || !name) { | |
| return; | |
| } | |
| let images = []; | |
| if (Array.isArray(risuData.additionalAssets)) { | |
| images = images.concat(risuData.additionalAssets); | |
| } | |
| if (Array.isArray(risuData.emotions)) { | |
| images = images.concat(risuData.emotions); | |
| } | |
| // No sprites to import | |
| if (images.length === 0) { | |
| return; | |
| } | |
| // Create sprites folder if it doesn't exist | |
| const spritesPath = path.join(directories.characters, name); | |
| if (!fs.existsSync(spritesPath)) { | |
| fs.mkdirSync(spritesPath); | |
| } | |
| // Path to sprites is not a directory. This should never happen. | |
| if (!fs.statSync(spritesPath).isDirectory()) { | |
| return; | |
| } | |
| console.log(`RisuAI: Found ${images.length} sprites for ${name}. Writing to disk.`); | |
| const files = fs.readdirSync(spritesPath); | |
| outer: for (const [label, fileBase64] of images) { | |
| // Remove existing sprite with the same label | |
| for (const file of files) { | |
| if (path.parse(file).name === label) { | |
| console.log(`RisuAI: The sprite ${label} for ${name} already exists. Skipping.`); | |
| continue outer; | |
| } | |
| } | |
| const filename = label + '.png'; | |
| const pathToFile = path.join(spritesPath, filename); | |
| writeFileAtomicSync(pathToFile, fileBase64, { encoding: 'base64' }); | |
| } | |
| // Remove additionalAssets and emotions from data (they are now in the sprites folder) | |
| delete data.data.extensions.risuai.additionalAssets; | |
| delete data.data.extensions.risuai.emotions; | |
| } catch (error) { | |
| console.error(error); | |
| } | |
| } | |
| const router = express.Router(); | |
| router.get('/get', jsonParser, function (request, response) { | |
| const name = String(request.query.name); | |
| const isSubfolder = name.includes('/'); | |
| const spritesPath = getSpritesPath(request.user.directories, name, isSubfolder); | |
| let sprites = []; | |
| try { | |
| if (spritesPath && fs.existsSync(spritesPath) && fs.statSync(spritesPath).isDirectory()) { | |
| sprites = fs.readdirSync(spritesPath) | |
| .filter(file => { | |
| const mimeType = mime.lookup(file); | |
| return mimeType && mimeType.startsWith('image/'); | |
| }) | |
| .map((file) => { | |
| const pathToSprite = path.join(spritesPath, file); | |
| return { | |
| label: path.parse(pathToSprite).name.toLowerCase(), | |
| path: `/characters/${name}/${file}`, | |
| }; | |
| }); | |
| } | |
| } | |
| catch (err) { | |
| console.log(err); | |
| } | |
| return response.send(sprites); | |
| }); | |
| router.post('/delete', jsonParser, async (request, response) => { | |
| const label = request.body.label; | |
| const name = request.body.name; | |
| if (!label || !name) { | |
| return response.sendStatus(400); | |
| } | |
| try { | |
| const spritesPath = path.join(request.user.directories.characters, name); | |
| // No sprites folder exists, or not a directory | |
| if (!fs.existsSync(spritesPath) || !fs.statSync(spritesPath).isDirectory()) { | |
| return response.sendStatus(404); | |
| } | |
| const files = fs.readdirSync(spritesPath); | |
| // Remove existing sprite with the same label | |
| for (const file of files) { | |
| if (path.parse(file).name === label) { | |
| fs.rmSync(path.join(spritesPath, file)); | |
| } | |
| } | |
| return response.sendStatus(200); | |
| } catch (error) { | |
| console.error(error); | |
| return response.sendStatus(500); | |
| } | |
| }); | |
| router.post('/upload-zip', urlencodedParser, async (request, response) => { | |
| const file = request.file; | |
| const name = request.body.name; | |
| if (!file || !name) { | |
| return response.sendStatus(400); | |
| } | |
| try { | |
| const spritesPath = path.join(request.user.directories.characters, name); | |
| // Create sprites folder if it doesn't exist | |
| if (!fs.existsSync(spritesPath)) { | |
| fs.mkdirSync(spritesPath); | |
| } | |
| // Path to sprites is not a directory. This should never happen. | |
| if (!fs.statSync(spritesPath).isDirectory()) { | |
| return response.sendStatus(404); | |
| } | |
| const spritePackPath = path.join(file.destination, file.filename); | |
| const sprites = await getImageBuffers(spritePackPath); | |
| const files = fs.readdirSync(spritesPath); | |
| for (const [filename, buffer] of sprites) { | |
| // Remove existing sprite with the same label | |
| const existingFile = files.find(file => path.parse(file).name === path.parse(filename).name); | |
| if (existingFile) { | |
| fs.rmSync(path.join(spritesPath, existingFile)); | |
| } | |
| // Write sprite buffer to disk | |
| const pathToSprite = path.join(spritesPath, filename); | |
| writeFileAtomicSync(pathToSprite, buffer); | |
| } | |
| // Remove uploaded ZIP file | |
| fs.rmSync(spritePackPath); | |
| return response.send({ count: sprites.length }); | |
| } catch (error) { | |
| console.error(error); | |
| return response.sendStatus(500); | |
| } | |
| }); | |
| router.post('/upload', urlencodedParser, async (request, response) => { | |
| const file = request.file; | |
| const label = request.body.label; | |
| const name = request.body.name; | |
| if (!file || !label || !name) { | |
| return response.sendStatus(400); | |
| } | |
| try { | |
| const spritesPath = path.join(request.user.directories.characters, name); | |
| // Create sprites folder if it doesn't exist | |
| if (!fs.existsSync(spritesPath)) { | |
| fs.mkdirSync(spritesPath); | |
| } | |
| // Path to sprites is not a directory. This should never happen. | |
| if (!fs.statSync(spritesPath).isDirectory()) { | |
| return response.sendStatus(404); | |
| } | |
| const files = fs.readdirSync(spritesPath); | |
| // Remove existing sprite with the same label | |
| for (const file of files) { | |
| if (path.parse(file).name === label) { | |
| fs.rmSync(path.join(spritesPath, file)); | |
| } | |
| } | |
| const filename = label + path.parse(file.originalname).ext; | |
| const spritePath = path.join(file.destination, file.filename); | |
| const pathToFile = path.join(spritesPath, filename); | |
| // Copy uploaded file to sprites folder | |
| fs.cpSync(spritePath, pathToFile); | |
| // Remove uploaded file | |
| fs.rmSync(spritePath); | |
| return response.sendStatus(200); | |
| } catch (error) { | |
| console.error(error); | |
| return response.sendStatus(500); | |
| } | |
| }); | |
| module.exports = { | |
| router, | |
| importRisuSprites, | |
| }; | |