From 9ef244fc3fcb30e2f2a38be82609d2d6dd64c9c9 Mon Sep 17 00:00:00 2001
From: Supositware <sup@libtar.de>
Date: Tue, 4 Apr 2023 16:06:31 +0000
Subject: [PATCH] Added todo, updated image fetch and added optout of sharing
 the end result

---
 commands/AI/txt2img.js | 21 +++++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/commands/AI/txt2img.js b/commands/AI/txt2img.js
index 625e391..9b9b97e 100644
--- a/commands/AI/txt2img.js
+++ b/commands/AI/txt2img.js
@@ -1,8 +1,16 @@
+/* TODO
+ *
+ * To be merged with commands/AI/img2img.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;
 
@@ -37,8 +45,15 @@ async function generate(i, prompt, client) {
 		steps: 50,
 		nsfw: i.channel.nsfw ? true : false,
 		censor_nsfw: i.channel.nsfw ? true : false,
+		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),
@@ -69,8 +84,10 @@ async function generate(i, prompt, client) {
 			let creditResponse = await fetch(`https://stablehorde.net/api/v2/users/${stableHordeID}`);
 			creditResponse = await creditResponse.json();
 
-			await fetch(checkResult.image)
-				.then(res => res.body.pipe(fs.createWriteStream(`${os.tmpdir()}/${i.id}.webp`)));
+			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`);