skynet/skynet_bot/frontend/telegram.py

#!/usr/bin/python

import logging

from datetime import datetime

import pynng

from telebot.async_telebot import AsyncTeleBot
from trio_asyncio import aio_as_trio

from ..constants import *

from . import *


PREFIX = 'tg'


async def run_skynet_telegram(tg_token: str):

    logging.basicConfig(level=logging.INFO)
    bot = AsyncTeleBot(tg_token)

    with open_skynet_rpc() as rpc_sock:

        async def _rpc_call(
            uid: int,
            method: str,
            params: dict
        ):
            return await rpc_call(
                rpc_sock, f'{PREFIX}+{uid}', method, params)

        @bot.message_handler(commands=['help'])
        async def send_help(message):
            await bot.reply_to(message, HELP_TEXT)

        @bot.message_handler(commands=['cool'])
        async def send_cool_words(message):
            await bot.reply_to(message, '\n'.join(COOL_WORDS))

        @bot.message_handler(commands=['txt2img'])
        async def send_txt2img(message):
            resp = await _rpc_call(
                message.from_user.id,
                'txt2img',
                {}
            )

        @bot.message_handler(commands=['redo'])
        async def redo_txt2img(message):
            resp = await _rpc_call(
                message.from_user.id,
                'redo',
                {}
            )

        @bot.message_handler(commands=['config'])
        async def set_config(message):
            params = message.text.split(' ')

            rpc_params = {}

            if len(params) < 3:
                bot.reply_to(message, 'wrong msg format')

            else:

                try:
                    attr = params[1]

                    if attr == 'algo':
                        val = params[2]
                        if val not in ALGOS:
                            raise ConfigUnknownAlgorithm

                    elif attr == 'step':
                        val = int(params[2])
                        val = max(min(val, MAX_STEP), MIN_STEP)

                    elif attr  == 'width':
                        val = max(min(int(params[2]), MAX_WIDTH), 16)
                        if val % 8 != 0:
                            raise ConfigSizeDivisionByEight

                    elif attr  == 'height':
                        val = max(min(int(params[2]), MAX_HEIGHT), 16)
                        if val % 8 != 0:
                            raise ConfigSizeDivisionByEight

                    elif attr == 'seed':
                        val = params[2]
                        if val == 'auto':
                            val = None
                        else:
                            val = int(params[2])

                    elif attr == 'guidance':
                        val = float(params[2])
                        val = max(min(val, MAX_GUIDANCE), 0)

                    elif attr == 'upscaler':
                        val = params[2]
                        if val == 'off':
                            val = None
                        elif val != 'x4':
                            raise ConfigUnknownUpscaler

                    else:
                        raise ConfigUnknownAttribute

                    resp = await _rpc_call(
                        message.from_user.id,
                        'config', {'attr': attr, 'val': val})

                    reply_txt = f'config updated! {attr} to {val}'

                except ConfigUnknownAlgorithm:
                    reply_txt = f'no algo named {val}'

                except ConfigUnknownAttribute:
                    reply_txt = f'\"{attr}\" not a configurable parameter'

                except ConfigUnknownUpscaler:
                    reply_txt = f'\"{val}\" is not a valid upscaler'

                except ConfigSizeDivisionByEight:
                    reply_txt = 'size must be divisible by 8!'

                except ValueError:
                    reply_txt = f'\"{val}\" is not a number silly'

                await bot.reply_to(message, reply_txt)

        @bot.message_handler(commands=['stats'])
        async def user_stats(message):
            resp = await _rpc_call(
                message.from_user.id,
                'stats',
                {}
            )
            stats = resp.result

            stats_str = f'generated: {stats["generated"]}\n'
            stats_str += f'joined: {stats["joined"]}\n'
            stats_str += f'role: {stats["role"]}\n'

            await bot.reply_to(
                message, stats_str)

        @bot.message_handler(commands=['donate'])
        async def donation_info(message):
            await bot.reply_to(
                message, DONATION_INFO)


        @bot.message_handler(func=lambda message: True)
        async def echo_message(message):
            if message.text[0] == '/':
                await bot.reply_to(message, UNKNOWN_CMD_TEXT)


        await aio_as_trio(bot.infinity_polling())