diff --git a/commands/AI/img2img b/commands/AI/img2img deleted file mode 100644 index bd837d8..0000000 --- a/commands/AI/img2img +++ /dev/null @@ -1,76 +0,0 @@ -import { SlashCommandBuilder, EmbedBuilder } from 'discord.js'; -import fetch from 'node-fetch'; -import fs from 'node:fs'; -import os from 'node:os'; - -const { stableHordeApi, stableHordeID } = process.env; - -export default { - data: new SlashCommandBuilder() - .setName('img2img') - .setDescription('AI generated image with stable diffusion (If credit are low it may be slow)') - .addAttachmentOption(option => - option.setName('image') - .setDescription('Image you want to modify') - .setRequired(true)) - .addStringOption(option => - option.setName('prompt') - .setDescription('What do you want the AI to generate?') - .setRequired(true)), - category: 'AI', - async execute(interaction, args, client) { - await interaction.deferReply(); - fetch(args.image.url) - .then((res) => { - const dest = fs.createWriteStream(`${os.tmpdir()}/${args.image.name}`); - res.body.pipe(dest); - dest.on('finish', () => { - const b64Image = fs.readFileSync(`${os.tmpdir()}/${args.image.name}`, { encoding: 'base64' }); - generate(interaction, args.prompt, b64Image, client); - }); - }); - }, -}; - -async function generate(i, prompt, b64Img) { - const body = { - prompt: prompt, - params: { - n: 1, - width: 512, - height: 512, - }, - cfg_scale: 9, - 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, - source_image: b64Img, - }; - - const fetchParameters = { - method: 'post', - body: JSON.stringify(body), - headers: { 'Content-Type': 'application/json', 'apikey': stableHordeApi }, - }; - - let response = await fetch('https://stablehorde.net/api/v2/generate/sync', fetchParameters); - - response = await response.json(); - - let creditResponse = await fetch(`https://stablehorde.net/api/v2/users/${stableHordeID}`); - creditResponse = await creditResponse.json(); - - const stableEmbed = new EmbedBuilder() - .setColor(i.member ? i.member.displayHexColor : 'Navy') - .setTitle(prompt) - .setURL('https://aqualxx.github.io/stable-ui/') - .setFooter({ text: `**Credit left: ${creditResponse.kudos}** Seed: ${response.generations[0].seed} worker ID: ${response.generations[0].worker_id} worker name: ${response.generations[0].worker_name}` }); - - fs.writeFileSync(`${os.tmpdir()}/${i.id}.png`, response.generations[0].img, 'base64'); - - await i.editReply({ embeds: [stableEmbed], files: [`${os.tmpdir()}/${i.id}.png`] }); -} diff --git a/commands/AI/img2img.js b/commands/AI/img2img.js new file mode 100644 index 0000000..9fdea40 --- /dev/null +++ b/commands/AI/img2img.js @@ -0,0 +1,153 @@ +/* TODO + * + * To be merged with commands/AI/txt2img.js + * +*/ +import { SlashCommandBuilder, EmbedBuilder, AttachmentBuilder, ButtonBuilder, ActionRowBuilder, ButtonStyle } from 'discord.js'; +import fetch from 'node-fetch'; +import os from 'node:os'; +import fs from 'node:fs'; +import stream from 'node:stream'; +import util from 'node:util'; + +import db from '../../models/index.js'; + +const { stableHordeApi, stableHordeID } = process.env; + +export default { + data: new SlashCommandBuilder() + .setName('img2img') + .setDescription('AI generated image with stable diffusion (If credit are low it may be slow)') + .addAttachmentOption(option => + option.setName('image') + .setDescription('Image you want to modify') + .setRequired(true)) + .addStringOption(option => + option.setName('prompt') + .setDescription('What do you want the AI to generate?') + .setRequired(true)), + category: 'AI', + async execute(interaction, args, client) { + await interaction.deferReply(); + + const streamPipeline = util.promisify(stream.pipeline); + const res = await fetch(args.image.url); + if (!res.ok) return interaction.editReply('An error has occured while trying to download your image.'); + await streamPipeline(res.body, fs.createWriteStream(`${os.tmpdir()}/${args.image.name}.webp`)); + + const b64Image = fs.readFileSync(`${os.tmpdir()}/${args.image.name}`, { encoding: 'base64' }); + generate(interaction, args.prompt, client, b64Image); + }, +}; + +async function generate(i, prompt, client, b64Img) { + console.log('Generating image'); + const body = { + prompt: prompt, + params: { + n: 1, + width: 512, + height: 512, + }, + cfg_scale: 9, + 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, + source_image: b64Img, + source_processing: 'img2img', + shared: true, + }; + + const isOptOut = await db.optout.findOne({ where: { userID: i.user.id } }); + + if (isOptOut) { + body.shared = false; + } + + const fetchParameters = { + method: 'post', + body: JSON.stringify(body), + headers: { 'Content-Type': 'application/json', 'apikey': stableHordeApi }, + }; + + let response = await fetch('https://stablehorde.net/api/v2/generate/async', fetchParameters); + + response = await response.json(); + let wait_time = 5000; + let checkURL = `https://stablehorde.net/api/v2/generate/check/${response.id}`; + const checking = setInterval(async () => { + const checkResult = await checkGeneration(checkURL); + + if (checkResult === undefined) return; + if (!checkResult.done) { + if (checkResult.wait_time < 0) { + console.log(checkResult); + 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; + } + 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`)); + + const generatedImg = new AttachmentBuilder(`${os.tmpdir()}/${i.id}.webp`); + + const stableEmbed = new EmbedBuilder() + .setColor(i.member ? i.member.displayHexColor : 'Navy') + .setTitle(prompt) + .setURL('https://aqualxx.github.io/stable-ui/') + .setImage(`attachment://${i.id}.webp`) + .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() + .setCustomId(`regenerate${i.user.id}`) + .setLabel('🔄 Regenerate') + .setStyle(ButtonStyle.Primary), + ); + + await i.editReply({ embeds: [stableEmbed], components: [row], files: [generatedImg] }); + + client.once('interactionCreate', async (interactionMenu) => { + if (i.user !== interactionMenu.user) return; + if (!interactionMenu.isButton) return; + if (interactionMenu.customId === `regenerate${interactionMenu.user.id}`) { + await interactionMenu.deferReply(); + await generate(interactionMenu, prompt, client); + } + }); + } + }, 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 }; + } +}