jack
e9f7b6e28e
Some checks failed
Build and Publish Docker Image / build_and_push (push) Failing after 19s
114 lines
3.4 KiB
TypeScript
114 lines
3.4 KiB
TypeScript
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-1.1-pro-ultra',
|
|
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.jpg',
|
|
});
|
|
|
|
// 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.');
|
|
}
|
|
}
|