Spaghetti Command

commands/edit-image/spaghetti.js

@@ -0,0 +1,41 @@
+const Command = require('../../framework/Command');
+const { PermissionFlagsBits } = require('discord.js');
+const request = require('node-superfetch');
+const { readFile } = require('fs/promises');
+const path = require('path');
+const { reactIfAble } = require('../../util/Util');
+const { LOADING_EMOJI_ID, SUCCESS_EMOJI_ID } = process.env;
+
+module.exports = class SpaghettiCommand extends Command {
+	constructor(client) {
+		super(client, {
+			name: 'spaghetti',
+			aliases: ['spaghettify', 'spaghet', 'pasta'],
+			group: 'edit-image',
+			memberName: 'spaghetti',
+			description: 'Draws an image or a user\'s avatar but as spaghetti.',
+			throttling: {
+				usages: 1,
+				duration: 120
+			},
+			clientPermissions: [PermissionFlagsBits.AttachFiles],
+			args: [
+				{
+					key: 'image',
+					type: 'image-or-avatar',
+					default: msg => msg.author.displayAvatarURL({ extension: 'png', size: 512 })
+				}
+			]
+		});
+	}
+
+	async run(msg, { image }) {
+		const { body } = await request.get(image);
+		const style = await readFile(path.join(__dirname, '..', '..', 'assets', 'images', 'spaghetti.jpg'));
+		await reactIfAble(msg, msg.author, LOADING_EMOJI_ID, '💬');
+		const attachment = await this.client.tensorflow.stylizeImage(body, style);
+		await reactIfAble(msg, msg.author, SUCCESS_EMOJI_ID, '✅');
+		if (Buffer.byteLength(attachment) > 2.5e+7) return msg.reply('Resulting image was above 25 MB.');
+		return msg.say({ files: [{ attachment, name: 'spaghetti.jpg' }] });
+	}
+};

package.json

@@ -1,6 +1,6 @@
 {
 	"name": "xiao",
-	"version": "150.1.0",
+	"version": "151.0.0",
 	"description": "Your personal server companion.",
 	"main": "Xiao.js",
 	"private": true,

structures/Tensorflow.js

@@ -1,7 +1,8 @@
-const tfnode = require('@tensorflow/tfjs-node');
+const tf = require('@tensorflow/tfjs-node');
 const nsfw = require('nsfwjs');
 const faceDetection = require('@tensorflow-models/face-detection');
 const faceModel = faceDetection.SupportedModels.MediaPipeFaceDetector;
+const path = require('path');
 
 module.exports = class Tensorflow {
 	constructor(client) {
@@ -9,6 +10,8 @@ module.exports = class Tensorflow {
 
 		this.nsfwjs = null;
 		this.faceDetector = null;
+		this.styleModel = null;
+		this.transformerModel = null;
 	}
 
 	async loadNSFWJS() {
@@ -23,20 +26,34 @@ module.exports = class Tensorflow {
 		return this.faceDetector;
 	}
 
+	async loadStyleModel() {
+		const model = await tf.loadGraphModel(path.join(__dirname, '..', 'tf_models', 'style_js', 'model.json'));
+		this.styleModel = model;
+		return this.styleModel;
+	}
+
+	async loadTransformerModel() {
+		const model = await tf.loadGraphModel(
+			path.join(__dirname, '..', 'tf_models', 'transformer_separable_js', 'model.json')
+		);
+		this.transformerModel = model;
+		return this.transformerModel;
+	}
+
 	async detectFaces(imgData) {
 		if (Buffer.byteLength(imgData) >= 8e+6) return 'size';
-		tfnode.setBackend('tensorflow');
-		const image = tfnode.node.decodeImage(imgData, 3);
-		tfnode.setBackend('cpu');
+		tf.setBackend('tensorflow');
+		const image = tf.node.decodeImage(imgData, 3);
+		tf.setBackend('cpu');
 		const faces = await this.faceDetector.estimateFaces(image);
-		tfnode.setBackend('tensorflow');
+		tf.setBackend('tensorflow');
 		image.dispose();
 		if (!faces || !faces.length) return null;
 		return faces;
 	}
 
 	async isImageNSFW(image, bool = true) {
-		const img = await tfnode.node.decodeImage(image, 3);
+		const img = await tf.node.decodeImage(image, 3);
 		const predictions = await this.nsfwjs.classify(img);
 		img.dispose();
 		if (bool) {
@@ -50,4 +67,21 @@ module.exports = class Tensorflow {
 		}
 		return predictions;
 	}
+
+	async stylizeImage(image, styleImg) {
+		const imageTensor = await tf.node.decodeImage(image, 3);
+		const loadedImage = imageTensor.div(tf.scalar(255)).expandDims();
+		imageTensor.dispose();
+		const styleTensor = tf.node.decodeImage(styleImg, 3);
+		const loadedStyle = styleTensor.div(tf.scalar(255)).expandDims();
+		styleTensor.dispose();
+		const stylePrediction = await this.styleModel.predict(loadedStyle);
+		loadedStyle.dispose();
+		const stylizedImage = await this.transformerModel.predict([loadedImage, stylePrediction.squeeze()]);
+		loadedImage.dispose();
+		stylePrediction.dispose();
+		const buffer = await tf.node.encodePng(stylizedImage.squeeze());
+		stylizedImage.dispose();
+		return buffer;
+	}
 };