Fix txt2img AI generation

commands/AI/txt2img.js

@@ -1,7 +1,5 @@
 import { SlashCommandBuilder, EmbedBuilder, ButtonBuilder, ActionRowBuilder, ButtonStyle } from 'discord.js';
 import fetch from 'node-fetch';
-import fs from 'node:fs';
-import os from 'node:os';
 
 const { stableHordeApi, stableHordeID } = process.env;
 
@@ -44,18 +42,41 @@ async function generate(i, prompt, client) {
 		headers: { 'Content-Type': 'application/json', 'apikey': stableHordeApi },
 	};
 
-	let response = await fetch('https://stablehorde.net/api/v2/generate/sync', fetchParameters);
+	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) {
+				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 imageData = await fetch(checkResult.image);
+			let imgBuffer = await imageData.arrayBuffer();
+			imgBuffer = Buffer.from(imgBuffer).toString('base64');
+			const img = `data:image/${imageData.headers.get('content-type')};base64,${imgBuffer}`;
+
 			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}` });
+				.setImage(`attachment:${img}`)
+				.setFooter({ text: `**Credit left: ${creditResponse.kudos}** Seed: ${checkResult.seed} worker ID: ${checkResult.worker_id} worker name: ${checkResult.worker_name}` });
 
 			const row = new ActionRowBuilder()
 				.addComponents(
@@ -65,9 +86,7 @@ async function generate(i, prompt, client) {
 						.setStyle(ButtonStyle.Primary),
 				);
 
-	fs.writeFileSync(`${os.tmpdir()}/${i.id}.png`, response.generations[0].img, 'base64');
+			await i.editReply({ embeds: [stableEmbed], components: [row] });
-
-	await i.editReply({ embeds: [stableEmbed], components: [row], files: [`${os.tmpdir()}/${i.id}.png`] });
 
 			client.once('interactionCreate', async (interactionMenu) => {
 				if (i.user !== interactionMenu.user) return;
@@ -77,4 +96,23 @@ async function generate(i, prompt, client) {
 					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 };
+	}
 }