Discordに投稿したらStable Diffusionで画像生成するプログラム

# はじめに

ローカル環境で動かしているStable DiffusionとDiscordを連携させるプログラムです。
Stable Diffusionの環境が無い友達ともワイワイやりながら画像生成ができるようにするために作りました。

また毎回英語にするのが面倒だったため、「discordに日本語で投稿->deeplで翻訳->Stable Diffusionで画像生成->画像をdiscordに投稿」までを一気にできるようにしました。
こんな感じで使えます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/114063/343dd37e-8cc0-df4a-7cba-ccbcb8b8ee06.png)

## 事前準備

* [Stable Diffusionの登録](https://huggingface.co/blog/stable_diffusion)
* Discord botの登録
* [Deepl（無料版）の登録](https://www.deepl.com/ja/pro/change-plan#developer)

を済ませておいてください。それぞれのtokenを使います。

また、Discordとdeeplは以下のライブラリを使っていますので各自インストールしてください。

```bash
pip install discord.py deepl
```

### 動作確認環境

* OS:Window11
* GPU:GTX1070Ti
* Python 3.10

## ソースコード

Stable Diffusion、discord、deeplそれぞれで取得したtokenをソースコード内に記載して実行してください。
「@bot名 <書いてほしい内容>」で画像生成してくれます。
「@bot名 残り」でdeeplの無料枠の残り文字数を確認できます。

```python
import datetime
import re

import deepl
import discord
import torch
from diffusers import StableDiffusionPipeline
from torch import autocast

deepl_token = <deeplのtoken>
discord_token = <Discordのtoken>
token = <Stable Diffusionのtoken>


def get_translation(text: str) -> str:
    result = translator.translate_text(text, target_lang="EN-US")
    return result.text  # type: ignore


def get_usage() -> str:
    usage = translator.get_usage()
    msg = ""
    if usage.any_limit_reached:
        msg += "Translation limit reached."
    if usage.character.valid:
        msg += f"Character usage: {usage.character.count} of {usage.character.limit}"
    if usage.document.valid:
        msg += f"Document usage: {usage.document.count} of {usage.document.limit}"
    return msg


class StableDiffusion:
    def __init__(self, token) -> None:
        model_id = "CompVis/stable-diffusion-v1-4"
        self.pipe = StableDiffusionPipeline.from_pretrained(model_id, revision="fp16", torch_dtype=torch.float16, use_auth_token=token)

    def generate_image(self, text: str) -> str:
        pipe = self.pipe.to("cuda")
        with autocast("cuda"):
            image = pipe(text)["sample"][0]
        now = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
        image.save(f"out/image_{now}.png")
        img_path = f"out/image_{now}.png"
        return img_path


if __name__ == "__main__":
    translator = deepl.Translator(deepl_token)
    intents = discord.Intents.default()
    intents.message_content = True
    client = discord.Client(intents=intents)
    sd = StableDiffusion(token)

    @client.event
    async def on_ready():
        print(f"We have logged in as {client.user}")

    @client.event
    async def on_message(message):
        if message.author.bot:
            return
        elif client.user in message.mentions:
            file = None
            t = re.sub(r"<@.*> ", "", message.content)
            if t == "残り":
                reply = get_usage()
            else:
                await message.channel.send(f"{message.author.mention} ちょっとまっててね。")
                print(f"Start generate image. {t}")
                text = get_translation(t)
                reply = f"{message.author.mention} {text}"
                file = discord.File(sd.generate_image(text))
            await message.channel.send(reply, file=file)

    client.run(discord_token)
```