Convert to Typescript

src/commands/image.ts

@@ -0,0 +1,113 @@
+import {
+  SlashCommandBuilder,
+  AttachmentBuilder,
+  ChatInputCommandInteraction,
+} from 'discord.js';
+import axios from 'axios';
+import Replicate, { Prediction } from 'replicate';
+import config from '../config';
+
+const replicate = new Replicate({
+  auth: config.replicateApiKey,
+});
+
+// Initialise the command data.
+export const data = new SlashCommandBuilder()
+  .setName('image')
+  .setDescription('Generate an image based on a prompt.')
+  .addStringOption((option) =>
+    option
+      .setName('prompt')
+      .setDescription('The prompt to generate an image from')
+      .setRequired(true)
+  );
+
+console.log(`Loaded ${data.name} command.`);
+
+/**
+ * Generate an image based on a prompt and send it back to the user.
+ * @param interaction The interaction that triggered the command.
+ * @returns A promise that resolves when the command is finished executing.
+ */
+export async function execute(
+  interaction: ChatInputCommandInteraction
+): Promise<void> {
+  await interaction.deferReply();
+
+  const prompt = interaction.options.get('prompt')?.value as string;
+
+  try {
+    // Create image generation prediction
+    const prediction = await replicate.predictions.create({
+      model: 'black-forest-labs/flux-schnell',
+      input: { prompt },
+    });
+
+    // Poll until the image generation is complete
+    const completedPrediction = await pollPredictionStatus(prediction.id);
+
+    if (!completedPrediction || !completedPrediction.output) {
+      throw new Error('Failed to generate the image.');
+    }
+
+    const imageUrl = completedPrediction.output[0];
+
+    // Download the generated image
+    const imageBuffer = await downloadImage(imageUrl);
+
+    // Create an attachment to send the image back to the user
+    const attachment = new AttachmentBuilder(imageBuffer, {
+      name: 'image.png',
+    });
+
+    // Edit the deferred reply to include the generated image
+    await interaction.editReply({ files: [attachment] });
+  } catch (error: any) {
+    // Provide a more informative error message to the user
+    console.error(error); // Log the error for debugging purposes
+    await interaction.editReply(`An error occurred: ${error.message}`);
+  }
+}
+
+/**
+ * Poll the status of a prediction until it is no longer in progress.
+ * @param predictionId The ID of the prediction to poll.
+ * @param maxAttempts The maximum number of attempts to poll the prediction.
+ * @param interval The interval between each polling attempt in milliseconds.
+ * @returns The final status of the prediction.
+ */
+async function pollPredictionStatus(
+  predictionId: string,
+  maxAttempts = 5,
+  interval = 2000
+): Promise<Prediction> {
+  for (let attempt = 0; attempt < maxAttempts; attempt++) {
+    const latestPrediction = await replicate.predictions.get(predictionId);
+
+    if (
+      latestPrediction.status !== 'starting' &&
+      latestPrediction.status !== 'processing'
+    ) {
+      return latestPrediction;
+    }
+
+    // Wait before checking again
+    await new Promise((resolve) => setTimeout(resolve, interval));
+  }
+
+  throw new Error('Prediction timed out.');
+}
+
+/**
+ * Download an image from a URL and return it as a Buffer.
+ * @param url The URL of the image to download.
+ * @returns A Buffer containing the downloaded image.
+ */
+async function downloadImage(url: string): Promise<Buffer> {
+  try {
+    const response = await axios.get(url, { responseType: 'arraybuffer' });
+    return Buffer.from(response.data);
+  } catch (error) {
+    throw new Error('Failed to download the image.');
+  }
+}