Haha-Yes/commands/AI/txt2img.js

144 lines
4.4 KiB
JavaScript
Raw Normal View History

/* TODO
*
* To be merged with commands/AI/img2img.js
*
*/
2023-02-20 23:05:21 +01:00
import { SlashCommandBuilder, EmbedBuilder, AttachmentBuilder, ButtonBuilder, ActionRowBuilder, ButtonStyle } from 'discord.js';
2022-10-16 20:09:48 +02:00
import fetch from 'node-fetch';
2023-02-20 23:05:21 +01:00
import os from 'node:os';
import fs from 'node:fs';
import stream from 'node:stream';
import util from 'node:util';
2023-02-20 23:05:21 +01:00
import db from '../../models/index.js';
2022-10-16 20:09:48 +02:00
2022-10-16 22:58:31 +02:00
const { stableHordeApi, stableHordeID } = process.env;
2022-10-16 22:30:00 +02:00
2022-10-16 20:09:48 +02:00
export default {
data: new SlashCommandBuilder()
2022-10-17 21:30:29 +02:00
.setName('txt2img')
2022-10-17 16:24:38 +02:00
.setDescription('AI generated image with stable diffusion (If credit are low it may be slow)')
2022-10-16 20:09:48 +02:00
.addStringOption(option =>
option.setName('prompt')
.setDescription('What do you want the AI to generate?')
.setRequired(true)),
2022-10-17 21:30:29 +02:00
category: 'AI',
2022-10-16 20:09:48 +02:00
async execute(interaction, args, client) {
await interaction.deferReply();
2022-10-16 22:37:34 +02:00
generate(interaction, args.prompt, client);
2022-10-16 20:09:48 +02:00
},
};
2022-10-16 22:37:34 +02:00
async function generate(i, prompt, client) {
2022-10-16 20:09:48 +02:00
const body = {
prompt: prompt,
params: {
n: 1,
width: 512,
height: 512,
},
2022-10-17 21:32:36 +02:00
cfg_scale: 9,
2022-10-17 16:24:38 +02:00
use_gfpgan: true,
use_real_esrgan: true,
use_ldsr: true,
use_upscaling: true,
steps: 50,
nsfw: i.channel.nsfw ? true : false,
censor_nsfw: i.channel.nsfw ? true : false,
shared: true,
2022-10-16 20:09:48 +02:00
};
const isOptOut = await db.optout.findOne({ where: { userID: i.user.id } });
if (isOptOut) {
body.shared = false;
}
2022-10-16 20:09:48 +02:00
const fetchParameters = {
method: 'post',
body: JSON.stringify(body),
2022-10-16 22:30:00 +02:00
headers: { 'Content-Type': 'application/json', 'apikey': stableHordeApi },
2022-10-16 20:09:48 +02:00
};
2023-02-20 22:28:01 +01:00
let response = await fetch('https://stablehorde.net/api/v2/generate/async', fetchParameters);
2022-10-16 20:09:48 +02:00
response = await response.json();
2023-02-20 22:28:01 +01:00
let wait_time = 5000;
let checkURL = `https://stablehorde.net/api/v2/generate/check/${response.id}`;
const checking = setInterval(async () => {
const checkResult = await checkGeneration(checkURL);
2022-10-16 22:58:31 +02:00
2023-02-20 22:28:01 +01:00
if (checkResult === undefined) return;
if (!checkResult.done) {
if (checkResult.wait_time < 0) {
clearInterval(checking);
return i.editReply({ content: 'No servers are currently available to fulfill your request, please try again later.' });
}
if (checkResult.wait_time === 0) {
checkURL = `https://stablehorde.net/api/v2/generate/status/${response.id}`;
}
wait_time = checkResult.wait_time;
2022-10-16 22:37:34 +02:00
}
2023-02-20 22:28:01 +01:00
else if (checkResult.done && checkResult.image) {
clearInterval(checking);
let creditResponse = await fetch(`https://stablehorde.net/api/v2/users/${stableHordeID}`);
creditResponse = await creditResponse.json();
const streamPipeline = util.promisify(stream.pipeline);
const res = await fetch(checkResult.image);
if (!res.ok) return i.editReply('An error has occured while trying to download your image.');
await streamPipeline(res.body, fs.createWriteStream(`${os.tmpdir()}/${i.id}.webp`));
2023-02-20 23:05:21 +01:00
const generatedImg = new AttachmentBuilder(`${os.tmpdir()}/${i.id}.webp`);
2023-02-20 22:28:01 +01:00
const stableEmbed = new EmbedBuilder()
.setColor(i.member ? i.member.displayHexColor : 'Navy')
.setTitle(prompt)
.setURL('https://aqualxx.github.io/stable-ui/')
2023-02-20 23:05:21 +01:00
.setImage(`attachment://${i.id}.webp`)
2023-02-20 22:28:01 +01:00
.setFooter({ text: `**Credit left: ${creditResponse.kudos}** Seed: ${checkResult.seed} worker ID: ${checkResult.worker_id} worker name: ${checkResult.worker_name}` });
const row = new ActionRowBuilder()
.addComponents(
new ButtonBuilder()
2023-04-05 18:12:56 +02:00
.setCustomId(`regenerate${i.user.id}${i.id}`)
2023-02-20 22:28:01 +01:00
.setLabel('🔄 Regenerate')
.setStyle(ButtonStyle.Primary),
);
2023-02-20 23:05:21 +01:00
await i.editReply({ embeds: [stableEmbed], components: [row], files: [generatedImg] });
2023-02-20 22:28:01 +01:00
2023-04-05 18:12:56 +02:00
listenButton(client, i, prompt);
2023-02-20 22:28:01 +01:00
}
}, wait_time);
}
async function checkGeneration(url) {
let check = await fetch(url);
check = await check.json();
if (!check.is_possible) {
return { done: false, wait_time: -1 };
}
if (check.done) {
if (!check.generations) {
return { done: false, wait_time: check.wait_time * 1000 };
}
return { done: true, image: check.generations[0].img, seed: check.generations[0].seed, worker_id: check.generations[0].worker_id, worker_name: check.generations[0].worker_name };
}
2022-10-16 20:09:48 +02:00
}
2023-04-05 18:12:56 +02:00
async function listenButton(client, interaction, prompt) {
client.once('interactionCreate', async (interactionMenu) => {
if (!interactionMenu.isButton()) return;
await interactionMenu.update({ components: [] });
if (interactionMenu.customId === `regenerate${interactionMenu.user.id}${interaction.id}`) {
await interactionMenu.deferReply();
await generate(interactionMenu, prompt, client);
}
});
}