Discordで動くStable DiffusionによるAI画像生成BOTを作った話

import { makeDiscordBot, readyDiscordBot, syncSlashCommands } from './lib/discord';
import { getEnv } from './lib/env';
import { slashCommandHandler } from './handler/slashCommandHandler';

const { discordBotToken, discordAppID, discordGuildID } = getEnv();
const bot = makeDiscordBot();

// BOT準備
readyDiscordBot(bot);

// スラッシュコマンドの同期
syncSlashCommands(bot, discordAppID, discordGuildID, discordBotToken);

// スラッシュコマンド検出
bot.on('interactionCreate', async (interaction) => {
  if (interaction.isCommand()) {
    await slashCommandHandler(interaction);
  }
});

bot.login(discordBotToken);

import { SlashCommandBuilder } from 'discord.js';

export const buildSlashCommands = ():
ReturnType<typeof SlashCommandBuilder.prototype.toJSON>[] => {
  return [
    new SlashCommandBuilder()
      .setName('generateimage')
      .setDescription('AIで画像生成します。')
      .addStringOption((option) =>
        option
          .setName("prompt")
          .setDescription("呪文1")
          .setRequired(true)
      )
      .addStringOption((option) =>
        option
          .setName("negative_prompt")
          .setDescription("呪文2")
          .setRequired(true)
      )
      .toJSON(),
    new SlashCommandBuilder()
      .setName('regenetateimage')
      .setDescription('前回と同じプロンプトでAIが画像生成します。')
      .toJSON(),
  ];
};

export const syncSlashCommands = async (
  bot: Client<boolean>,
  clientId: string,
  guildId: string,
  token: string,
) => {
  try
  {
    const slashCommandBulders = buildSlashCommands();

    const rest = new REST({ version: '9' }).setToken(token);
    console.log('Started refreshing application (/) commands.');

    await rest.put(Routes.applicationGuildCommands(clientId, guildId), {
      body: slashCommandBulders,
    });

    console.log('Successfully reloaded application (/) commands.');
  } catch (error) {
    console.error(error);
  }
};

import { CommandInteraction } from 'discord.js';
import sdwebui from 'node-sd-webui';
import fs from 'fs';
import { ImgurClient } from "imgur"
import { getEnv } from '../lib/env';

interface Prompts {
  prompt :string,
  negativePrompt: string,
  width: number,
  height: number,
  steps: number,
};
const promts = new Map<string, Prompts>();
const STABLE_DIFFUSION_URL = 'http://127.0.0.1:7860';

const { imgurClientID, imgurClientSecret } = getEnv();
const imgurClient = new ImgurClient({
  clientId: imgurClientID,
  clientSecret: imgurClientSecret,
});

export const slashCommandHandler = async (interaction: CommandInteraction) => {
  const { commandName } = interaction;

  switch (commandName) {
    case 'generateimage':
      generateImageHandler(interaction);
      break;
    case "regenerateimage":
      getReGenerateImageHandler(interaction);
      break;
    default:
      break;
  }
};

const generateImageHandler = async (interaction: CommandInteraction) => {
  const prompt = interaction.options.get('prompt') ? 
    interaction.options.get('prompt')?.value as string :
    promts.get(interaction.user.id)?.prompt as string
  const negativePrompt = interaction.options.get('negative_prompt')?
    interaction.options.get('negative_prompt')?.value as string :
    promts.get(interaction.user.id)?.negativePrompt as string
  const width = interaction.options.get('width')?
    interaction.options.get('width')?.value as number :
    promts.get(interaction.user.id)?.width as number
  const height = interaction.options.get('height')?
    interaction.options.get('height')?.value as number :
    promts.get(interaction.user.id)?.height as number
  const steps = interaction.options.get('steps')?
    interaction.options.get('steps')?.value as number :
    promts.get(interaction.user.id)?.steps as number

  promts.set(interaction.user.id,{
    prompt: prompt,
    negativePrompt: negativePrompt,
    width: width,
    height: height,
    steps: steps,
  });
  
  const client = sdwebui({ apiUrl: STABLE_DIFFUSION_URL });

  try {
    await interaction.deferReply();
    const response = await client.txt2img({
      prompt: prompt,
      negativePrompt: negativePrompt,
      width: width,
      height: height,
      steps: steps,
      seed: -1,
      batchSize: 1,
      hires: {
        steps: 0,
        denoisingStrength: 0.7,
        upscaler: "Latent",
        upscaleBy: 2,
        resizeWidthTo: 1024,
        resizeHeigthTo: 1024,
      }
    });
    console.log('parameters', response.parameters);

    const info = JSON.parse(response.info);
    console.log('info', info);
    const fileName = `image-${info.job_timestamp}.png`;

    response.images.forEach((image) => {
      fs.writeFileSync(`./out/${fileName}`, image, 'base64');
    });

    const base64data = fs.readFileSync(`./out/${fileName}`, { encoding: "base64" })
    const imgurResponse = await imgurClient.upload({
      image: base64data,
      type: 'base64'
    });
    console.log(imgurResponse.data);

    await interaction.editReply(
      `${interaction.user.displayName}さんが画像生成しました。\n- ${prompt}\n- ${negativePrompt}\n- ${imgurResponse.data.link}`
    );
  } catch (error) {
    await interaction.editReply(
      `ごめんなさい！${interaction.user.displayName}さん！画像生成に失敗したよ🥺`
    );
    console.log(`${error}`);
  }
};

const getReGenerateImageHandler = async (interaction: CommandInteraction) => {
  if (promts.size !== 0) {
    generateImageHandler(interaction);
  } else {
    await interaction.reply(`前回のプロンプトが見つからなかったよ🥺`);
  }
};