はじめに

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

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

事前準備

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

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

pip install discord.py deepl

動作確認環境

OS:Window11
GPU:GTX1070Ti
Python 3.10

ソースコード

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

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)

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

はじめに

事前準備

動作確認環境

ソースコード