Updated image gen

This commit is contained in:
Atridad Lahiji 2023-11-10 18:36:44 -07:00
parent 8f77726ea1
commit 1ca5dbf987
No known key found for this signature in database
3 changed files with 43 additions and 225 deletions

View file

@ -1,93 +0,0 @@
import { ApplyOptions } from '@sapphire/decorators';
import { Args, BucketScope, Command } from '@sapphire/framework';
import { AttachmentBuilder, Message } from 'discord.js';
import OpenAI from 'openai';
const openai = new OpenAI({
apiKey: process.env.OPENAI_API_KEY
});
// @ts-ignore
@ApplyOptions<Command.Options>({
description: 'Shhhh!',
options: ['prompt'],
// 10mins
cooldownDelay: 400_000,
cooldownLimit: 1,
// Yes... I did hardcode myself.
cooldownFilteredUsers: ['83679718401904640'],
cooldownScope: BucketScope.User
})
export class UserCommand extends Command {
// Register Chat Input and Context Menu command
public override registerApplicationCommands(registry: Command.Registry) {
registry.registerChatInputCommand((builder) =>
builder
.setName(this.name)
.setDescription(this.description)
.addStringOption((option) =>
option.setName('prompt').setDescription('The prompt you will use to generate an image!').setRequired(true)
)
);
}
// Message command
public async messageRun(message: Message, args: Args) {
return this.pic(message, args.getOption('prompt') || 'Scold me for not passing any prompt in.');
}
// Chat Input (slash) command
public async chatInputRun(interaction: Command.ChatInputCommandInteraction) {
return this.pic(interaction, interaction.options.getString('prompt') || 'NOTHING');
}
private async pic(interactionOrMessage: Message | Command.ChatInputCommandInteraction | Command.ContextMenuCommandInteraction, prompt: string) {
const askMessage =
interactionOrMessage instanceof Message
? await interactionOrMessage.channel.send({ content: '🤔 Thinking... 🤔' })
: await interactionOrMessage.reply({ content: '🤔 Thinking... 🤔', fetchReply: true });
const response = await openai.images.generate({
model: 'dall-e-3',
prompt,
n: 1,
size: '1024x1024'
});
const imageUrl = response.data[0].url || '';
// get an array buffer
const imageBuffer = await fetch(imageUrl).then((r) => r.arrayBuffer());
if (response.data.length === 0) {
const content = `Sorry, I can't complete the prompt for: ${prompt}`;
if (interactionOrMessage instanceof Message) {
return askMessage.edit({ content });
}
return interactionOrMessage.editReply({
content: content
});
} else {
const imageAttachment: AttachmentBuilder[] = [];
imageAttachment.push(
new AttachmentBuilder(Buffer.from(new Uint8Array(imageBuffer)), {
name: 'response.jpg',
description: "Himbot's Response"
})
);
const content = `Prompt: ${prompt}:`;
if (interactionOrMessage instanceof Message) {
return askMessage.edit({ content, files: imageAttachment });
}
return interactionOrMessage.editReply({
content,
files: imageAttachment
});
}
}
}

View file

@ -1,15 +1,18 @@
import { ApplyOptions } from '@sapphire/decorators'; import { ApplyOptions } from '@sapphire/decorators';
import { Args, BucketScope, Command } from '@sapphire/framework'; import { Args, BucketScope, Command } from '@sapphire/framework';
import { AttachmentBuilder, Message } from 'discord.js'; import { AttachmentBuilder, Message } from 'discord.js';
import OpenAI from 'openai';
// This is literally the worlds messiest TS code. Please don't judge me... const openai = new OpenAI({
apiKey: process.env.OPENAI_API_KEY
});
// @ts-ignore // @ts-ignore
@ApplyOptions<Command.Options>({ @ApplyOptions<Command.Options>({
description: 'Make a picture... but high res!', description: 'Pic... but better!',
options: ['prompt', 'number of pictures'], options: ['prompt'],
// 10mins // 10mins
cooldownDelay: 300_000, cooldownDelay: 400_000,
cooldownLimit: 1, cooldownLimit: 1,
// Yes... I did hardcode myself. // Yes... I did hardcode myself.
cooldownFilteredUsers: ['83679718401904640'], cooldownFilteredUsers: ['83679718401904640'],
@ -25,85 +28,38 @@ export class UserCommand extends Command {
.addStringOption((option) => .addStringOption((option) =>
option.setName('prompt').setDescription('The prompt you will use to generate an image!').setRequired(true) option.setName('prompt').setDescription('The prompt you will use to generate an image!').setRequired(true)
) )
.addStringOption((option) =>
option
.setName('amount')
.setDescription('The number of images you would like to generate. Maximum 2.')
.setChoices(
...[
{ name: '1', value: '1' },
{ name: '2', value: '2' }
]
)
)
); );
} }
// Message command // Message command
public async messageRun(message: Message, args: Args) { public async messageRun(message: Message, args: Args) {
const amount = Math.abs(Number(args.getOption('amount'))); return this.pic(message, args.getOption('prompt') || 'NOTHING');
return this.picHr(message, args.getOption('prompt') || 'Scold me for not passing any prompt in.', amount <= 2 ? amount : 1);
} }
// Chat Input (slash) command // Chat Input (slash) command
public async chatInputRun(interaction: Command.ChatInputCommandInteraction) { public async chatInputRun(interaction: Command.ChatInputCommandInteraction) {
const amount = Number(interaction.options.getString('amount')); return this.pic(interaction, interaction.options.getString('prompt') || 'NOTHING');
return this.picHr(interaction, interaction.options.getString('prompt') || 'NOTHING', amount || 1);
} }
private async picHr( private async pic(interactionOrMessage: Message | Command.ChatInputCommandInteraction | Command.ContextMenuCommandInteraction, prompt: string) {
interactionOrMessage: Message | Command.ChatInputCommandInteraction | Command.ContextMenuCommandInteraction,
prompt: string,
amount: number
) {
const askMessage = const askMessage =
interactionOrMessage instanceof Message interactionOrMessage instanceof Message
? await interactionOrMessage.channel.send({ content: '🤔 Thinking... 🤔' }) ? await interactionOrMessage.channel.send({ content: '🤔 Thinking... 🤔' })
: await interactionOrMessage.reply({ content: '🤔 Thinking... 🤔', fetchReply: true }); : await interactionOrMessage.reply({ content: '🤔 Thinking... 🤔', fetchReply: true });
const creditCountResponse = await fetch(`https://api.stability.ai/v1/user/balance`, { try {
method: 'GET', const response = await openai.images.generate({
headers: { model: 'dall-e-3',
'Content-Type': 'application/json', prompt,
Authorization: `Bearer ${process.env.STABILITY_API_KEY}` n: 1,
} size: '1024x1024'
});
const balance = (await creditCountResponse.json()).credits || 0;
if (balance > 5) {
const imageGenResponse = await fetch(`https://api.stability.ai/v1/generation/stable-diffusion-xl-1024-v1-0/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: 6,
clip_guidance_preset: 'FAST_BLUE',
height: 1024,
width: 1024,
samples: amount,
steps: 32,
seed: Number(String(interactionOrMessage.member?.user.id).substring(0, 5)) || 0
})
}); });
interface GenerationResponse { const imageUrl = response.data[0].url || '';
artifacts: Array<{ // get an array buffer
base64: string; const imageBuffer = await fetch(imageUrl).then((r) => r.arrayBuffer());
seed: number;
finishReason: string;
}>;
}
if (!imageGenResponse.ok) { if (response.data.length === 0) {
const content = `Sorry, I can't complete the prompt for: ${prompt}`; const content = `Sorry, I can't complete the prompt for: ${prompt}`;
if (interactionOrMessage instanceof Message) { if (interactionOrMessage instanceof Message) {
@ -114,34 +70,16 @@ export class UserCommand extends Command {
content: content content: content
}); });
} else { } else {
const responseJSON = (await imageGenResponse.json()) as GenerationResponse;
const imageAttachment: AttachmentBuilder[] = []; const imageAttachment: AttachmentBuilder[] = [];
for (let i = 0; i < responseJSON.artifacts.length; i++) { imageAttachment.push(
imageAttachment.push( new AttachmentBuilder(Buffer.from(new Uint8Array(imageBuffer)), {
new AttachmentBuilder(Buffer.from(responseJSON.artifacts[i].base64, 'base64'), { name: 'himbot_response.jpg',
name: 'response.jpg', description: `An image generated by Himbot using the prompt: ${prompt}`
description: "Himbot's Response" })
}) );
);
}
const newCreditCountResponse = await fetch(`https://api.stability.ai/v1/user/balance`, { const content = `Prompt: ${prompt}:`;
method: 'GET',
headers: {
'Content-Type': 'application/json',
Authorization: `Bearer ${process.env.STABILITY_API_KEY}`
}
});
const newBalance = (await newCreditCountResponse.json()).credits || 0;
const content =
`Credits Used: ${balance - newBalance}\nPrompt: ${prompt}${
balance <= 300
? `\n\n⚠I am now at ${balance} credits. If you'd like to help fund this command, please type "/support" for details!`
: ''
}` || 'ERROR!';
if (interactionOrMessage instanceof Message) { if (interactionOrMessage instanceof Message) {
return askMessage.edit({ content, files: imageAttachment }); return askMessage.edit({ content, files: imageAttachment });
@ -152,16 +90,14 @@ export class UserCommand extends Command {
files: imageAttachment files: imageAttachment
}); });
} }
} else { } catch (error) {
const content = `Oops! We're out of credits for this. If you'd like to help fund this command, please type "/support" for details!`; const content = "Sorry, I can't complete the prompt for: " + prompt + '\n' + 'Error: ' + error;
if (interactionOrMessage instanceof Message) { if (interactionOrMessage instanceof Message) {
return askMessage.edit({ content }); return askMessage.edit({ content });
} }
return interactionOrMessage.editReply({ return interactionOrMessage.editReply({ content });
content: content
});
} }
} }
} }

View file

@ -2,14 +2,14 @@ import { ApplyOptions } from '@sapphire/decorators';
import { Args, BucketScope, Command } from '@sapphire/framework'; import { Args, BucketScope, Command } from '@sapphire/framework';
import { AttachmentBuilder, Message } from 'discord.js'; import { AttachmentBuilder, Message } from 'discord.js';
// This is literally the worlds messiest TS code. Please don't judge me... // This is literally the world's messiest Typescript code. Please don't judge me...
// @ts-ignore // @ts-ignore
@ApplyOptions<Command.Options>({ @ApplyOptions<Command.Options>({
description: 'Make a picture!', description: 'Make a picture... but high res!',
options: ['prompt', 'amount'], options: ['prompt'],
// 10mins // 10mins
cooldownDelay: 100_000, cooldownDelay: 300_000,
cooldownLimit: 1, cooldownLimit: 1,
// Yes... I did hardcode myself. // Yes... I did hardcode myself.
cooldownFilteredUsers: ['83679718401904640'], cooldownFilteredUsers: ['83679718401904640'],
@ -25,39 +25,20 @@ export class UserCommand extends Command {
.addStringOption((option) => .addStringOption((option) =>
option.setName('prompt').setDescription('The prompt you will use to generate an image!').setRequired(true) option.setName('prompt').setDescription('The prompt you will use to generate an image!').setRequired(true)
) )
.addStringOption((option) =>
option
.setName('amount')
.setDescription('The number of images you would like to generate. Maximum 4.')
.setChoices(
...[
{ name: '1', value: '1' },
{ name: '2', value: '2' },
{ name: '3', value: '3' },
{ name: '4', value: '4' }
]
)
)
); );
} }
// Message command // Message command
public async messageRun(message: Message, args: Args) { public async messageRun(message: Message, args: Args) {
const amount = Math.abs(Number(args.getOption('amount'))); return this.picHr(message, args.getOption('prompt') || 'Scold me for not passing any prompt in.');
return this.pic(message, args.getOption('prompt') || 'Scold me for not passing any prompt in.', amount <= 4 ? amount : 1);
} }
// Chat Input (slash) command // Chat Input (slash) command
public async chatInputRun(interaction: Command.ChatInputCommandInteraction) { public async chatInputRun(interaction: Command.ChatInputCommandInteraction) {
const amount = Number(interaction.options.getString('amount')); return this.picHr(interaction, interaction.options.getString('prompt') || 'NOTHING');
return this.pic(interaction, interaction.options.getString('prompt') || 'NOTHING', amount || 1);
} }
private async pic( private async picHr(interactionOrMessage: Message | Command.ChatInputCommandInteraction | Command.ContextMenuCommandInteraction, prompt: string) {
interactionOrMessage: Message | Command.ChatInputCommandInteraction | Command.ContextMenuCommandInteraction,
prompt: string,
amount: number
) {
const askMessage = const askMessage =
interactionOrMessage instanceof Message interactionOrMessage instanceof Message
? await interactionOrMessage.channel.send({ content: '🤔 Thinking... 🤔' }) ? await interactionOrMessage.channel.send({ content: '🤔 Thinking... 🤔' })
@ -74,7 +55,7 @@ export class UserCommand extends Command {
const balance = (await creditCountResponse.json()).credits || 0; const balance = (await creditCountResponse.json()).credits || 0;
if (balance > 5) { if (balance > 5) {
const imageGenResponse = await fetch(`https://api.stability.ai/v1/generation/stable-diffusion-xl-beta-v2-2-2/text-to-image`, { const imageGenResponse = await fetch(`https://api.stability.ai/v1/generation/stable-diffusion-xl-1024-v1-0/text-to-image`, {
method: 'POST', method: 'POST',
headers: { headers: {
'Content-Type': 'application/json', 'Content-Type': 'application/json',
@ -89,9 +70,8 @@ export class UserCommand extends Command {
], ],
cfg_scale: 6, cfg_scale: 6,
clip_guidance_preset: 'FAST_BLUE', clip_guidance_preset: 'FAST_BLUE',
height: 512, height: 1024,
width: 512, width: 1024,
samples: amount,
steps: 32, steps: 32,
seed: Number(String(interactionOrMessage.member?.user.id).substring(0, 5)) || 0 seed: Number(String(interactionOrMessage.member?.user.id).substring(0, 5)) || 0
}) })
@ -122,8 +102,8 @@ export class UserCommand extends Command {
for (let i = 0; i < responseJSON.artifacts.length; i++) { for (let i = 0; i < responseJSON.artifacts.length; i++) {
imageAttachment.push( imageAttachment.push(
new AttachmentBuilder(Buffer.from(responseJSON.artifacts[i].base64, 'base64'), { new AttachmentBuilder(Buffer.from(responseJSON.artifacts[i].base64, 'base64'), {
name: 'response.jpg', name: 'himbot_response.jpg',
description: "Himbot's Response" description: `An image generated by Himbot using the prompt: ${prompt}`
}) })
); );
} }
@ -138,12 +118,7 @@ export class UserCommand extends Command {
const newBalance = (await newCreditCountResponse.json()).credits || 0; const newBalance = (await newCreditCountResponse.json()).credits || 0;
const content = const content = `Credits Used: ${balance - newBalance}\nPrompt: ${prompt}` || 'ERROR!';
`Credits Used: ${balance - newBalance}\nPrompt: ${prompt}${
balance <= 300
? `\n\n⚠I am now at ${balance} credits. If you'd like to help fund this command, please type "/support" for details!`
: ''
}` || 'ERROR!';
if (interactionOrMessage instanceof Message) { if (interactionOrMessage instanceof Message) {
return askMessage.edit({ content, files: imageAttachment }); return askMessage.edit({ content, files: imageAttachment });