arty/arty/bot.py

import asyncio
import random
import time
from io import BytesIO
from threading import Thread

import discord
import torch

from make_image import make_image


class Bot:
    SORRY = [
        'my server sucks',
        'my server is slow as fuck',
        "i'm running on a toaster",
        "i'm running on a flip-phone",
        "i'm running on a tamagochi",
        "i'm running on a potato",
        "i'm not exactly running on a nasa computer",
        'i lack the recources to do this fast',
        "i'm sleepy",
        "i'm lazy",
    ]
    NO_GPU = [
        'my server lacks a graphics card',
        "i'm running on a v-server with no gpu",
        'my creator is too poor to get a server with gpu',
        'i lack a gpu',
        "i can't run on cuda (no gpu)"
    ]

    def __init__(self, secret: str):
        self._client = None
        self._secret = secret

        self._setup()

    def run(self):
        # note: blocking!
        self._client.run(self._secret)

    async def on_ready(self):
        print(f'We have logged in as {self._client.user}')

    async def on_message(self, message):
        print(f'{message.author}: {message.clean_content} ({message.mentions})')
        if message.author == self._client.user:
            return

        if message.clean_content.startswith(self._msg_prefix):
            await self._draw_image(message)

    @property
    def _msg_prefix(self):
        return f'@{self._client.user.name} draw'

    def _setup(self):
        intents = discord.Intents.default()
        intents.message_content = True

        self._client = discord.Client(intents=intents)
        self._client.event(self.on_ready)
        self._client.event(self.on_message)

    async def _draw_image(self, message: discord.Message):
        await self._ack_prompt(message)

        prompt = message.clean_content[len(self._msg_prefix) + 1:]
        print(f'drawing an image, prompt="{prompt}"')

        image_grid = await self._make_image(prompt)

        buffer = BytesIO()
        image_grid['result'].convert('RGB')
        image_grid['result'].save(buffer, format='JPEG')

        buffer.seek(0)
        picture = discord.File(buffer, filename='result.jpeg')

        if image_grid['nsfw'][0]:
            await message.channel.send(
                "oh, this one's spicy (according to the nsfw filter, that i don't care about)",
            )

        await message.channel.send(f'{message.author.mention} {prompt}:', file=picture)

    async def _ack_prompt(self, message: discord.Message):
        messages = self.SORRY
        if not torch.cuda.is_available():
            messages += self.NO_GPU

        sorry = random.choice(messages)
        await message.channel.send(f'No problem. Just give me some time, {sorry}.')

    async def _make_image(self, prompt: str) -> dict:
        def start_background_loop(loop: asyncio.AbstractEventLoop) -> None:
            asyncio.set_event_loop(loop)
            loop.run_forever()

        async def task(_prompt) -> dict:
            return make_image(_prompt)

        loop = asyncio.new_event_loop()
        t0 = time.time()

        thread = Thread(target=start_background_loop, args=(loop,), daemon=True)
        thread.start()

        task = asyncio.run_coroutine_threadsafe(task(prompt), loop)
        result = task.result()

        loop.stop()
        print(f'prompt rendered in {time.time() - t0}s')
        return result