From 89c0ac4e5ebef57a7b77b3e84bf1a8824edeae1f Mon Sep 17 00:00:00 2001
From: Atridad Lahiji <88056492+atridadl@users.noreply.github.com>
Date: Fri, 23 Jun 2023 20:53:11 -0600
Subject: [PATCH] better pics

---
 src/.env.example    |  2 ++
 src/commands/pic.ts | 70 +++++++++++++++++++++++++++++++++------------
 2 files changed, 54 insertions(+), 18 deletions(-)

diff --git a/src/.env.example b/src/.env.example
index 92662fa..cb27056 100644
--- a/src/.env.example
+++ b/src/.env.example
@@ -1,2 +1,4 @@
 # Tokens
 DISCORD_TOKEN=
+OPENAI_API_KEY=
+STABILITY_API_KEY=
\ No newline at end of file
diff --git a/src/commands/pic.ts b/src/commands/pic.ts
index 19f8668..c90c1b5 100644
--- a/src/commands/pic.ts
+++ b/src/commands/pic.ts
@@ -1,12 +1,6 @@
 import { ApplyOptions } from '@sapphire/decorators';
 import { Args, BucketScope, Command } from '@sapphire/framework';
-import { Message } from 'discord.js';
-import { Configuration, OpenAIApi } from 'openai';
-
-const configuration = new Configuration({
-	apiKey: process.env.OPENAI_API_KEY
-});
-const openai = new OpenAIApi(configuration);
+import { AttachmentBuilder, Message } from 'discord.js';
 
 @ApplyOptions<Command.Options>({
 	description: 'Make a picture!',
@@ -45,20 +39,60 @@ export class UserCommand extends Command {
 				? await interactionOrMessage.channel.send({ content: '🤔 Thinking... 🤔' })
 				: await interactionOrMessage.reply({ content: '🤔 Thinking... 🤔', fetchReply: true });
 
-		const imageResponse = await openai.createImage({
-			prompt,
-			n: 1,
-			size: '512x512'
+		const response = await fetch(`https://api.stability.ai/v1/generation/stable-diffusion-xl-beta-v2-2-2/text-to-image`, {
+			method: 'POST',
+			headers: {
+				'Content-Type': 'application/json',
+				Accept: 'application/json',
+				Authorization: `Bearer ${process.env.STABILITY_API_KEY}`
+			},
+			body: JSON.stringify({
+				text_prompts: [
+					{
+						text: prompt
+					}
+				],
+				cfg_scale: 7,
+				clip_guidance_preset: 'FAST_BLUE',
+				height: 512,
+				width: 512,
+				samples: 1,
+				steps: 50
+			})
 		});
 
-		const content = `Prompt: ${prompt}\nResult: ${imageResponse.data.data[0].url}` || 'ERROR!';
-
-		if (interactionOrMessage instanceof Message) {
-			return askMessage.edit({ content });
+		interface GenerationResponse {
+			artifacts: Array<{
+				base64: string;
+				seed: number;
+				finishReason: string;
+			}>;
 		}
 
-		return interactionOrMessage.editReply({
-			content: content
-		});
+		if (!response.ok) {
+			const content = `Prompt: ${prompt}` || 'ERROR!';
+
+			if (interactionOrMessage instanceof Message) {
+				return askMessage.edit({ content });
+			}
+
+			return interactionOrMessage.editReply({
+				content: content
+			});
+		} else {
+			const responseJSON = (await response.json()) as GenerationResponse;
+			const imageAttachment = new AttachmentBuilder(Buffer.from(responseJSON.artifacts[0].base64, 'base64'));
+
+			const content = `Prompt: ${prompt}` || 'ERROR!';
+
+			if (interactionOrMessage instanceof Message) {
+				return askMessage.edit({ content, files: [imageAttachment] });
+			}
+
+			return interactionOrMessage.editReply({
+				content: content,
+				files: [imageAttachment]
+			});
+		}
 	}
 }