diff --git a/src/commands/imagine.js b/src/commands/imagine.js new file mode 100644 index 0000000..102b5b3 --- /dev/null +++ b/src/commands/imagine.js @@ -0,0 +1,76 @@ +import { SlashCommandBuilder } from "discord.js"; +import { WorkersAI } from "../util/models/index.js"; + +/** @type {import('./index.js').Command} */ +export default { + data: new SlashCommandBuilder() + .setName("imagine") + .setDescription("Generate an image.") + .addStringOption((o) => + o + .setName("model") + .setDescription("Enter a description of what you're generating.") + .setChoices([ + { + name: "lykon/dreamshaper-8-lcm", + value: "@cf/lykon/dreamshaper-8-lcm", + }, + { + name: "bytedance/stable-diffusion-xl-lightning", + value: "@cf/bytedance/stable-diffusion-xl-lightning", + }, + { + name: "stabilityai/stable-diffusion-xl-base-1.0", + value: "@cf/stabilityai/stable-diffusion-xl-base-1.0", + }, + ]) + .setRequired(true), + ) + .addStringOption((o) => + o.setName("prompt").setDescription("Enter a description of what you're generating.").setRequired(true), + ) + .toJSON(), + async execute(interaction) { + const workersAI = new WorkersAI({ + accountId: process.env.CLOUDFLARE_ACCOUNT_ID, + token: process.env.CLOUDFLARE_ACCOUNT_TOKEN, + }); + + await interaction.deferReply(); + const prompt = interaction.options.getString("prompt"); + const model = interaction.options.getString("model"); + + const callToModel = await (async () => { + const prefix = model?.split("/")?.[0]; + if (prefix !== "@cf") return; + + return await workersAI + .callModel( + { + model, + input: { + prompt, + }, + }, + true, + ) + .then((r) => r.arrayBuffer()) + .catch(() => (e) => { + console.error(e); + return null; + }); + })(); + + if (callToModel === null) + return await interaction.editReply({ + content: `The model did not generate an image.`, + }); + + const buffer = Buffer.from(callToModel); + + await interaction.editReply({ + content: `\`${prompt}\`\n*generated with \`${model}\`*`, + files: [buffer], + }); + }, +};