main.py

import logging
import pickle
from copy import copy
from datetime import datetime
from datetime import datetime as dt
from io import BytesIO
from os import getenv, remove
from os.path import exists, join
from random import randint, seed
from threading import Thread
from uuid import uuid4

import diskcache as dc
import pika
import sentry_sdk
import telegram
import webuiapi
from better_profanity import profanity
from deep_translator import GoogleTranslator
from nsfw_detector import predict
from PIL import Image
from pyairtable import Api, Base, Table
from sentry_sdk.consts import OP
from sentry_sdk.hub import Hub
from telegram import (InlineKeyboardButton, InlineKeyboardMarkup,
                      ReplyKeyboardMarkup, ReplyKeyboardRemove, Update)
from telegram import __version__ as TG_VER
from telegram.ext import (Application, CallbackQueryHandler, CommandHandler,
                          ContextTypes, MessageHandler, filters)
from tqdm import tqdm
from urllib.parse import urlparse
import default
from fastparquet import ParquetFile
# from telethon.sync import TelegramClient
# from telethon import connection

from low import lowpriority
lowpriority()

try:
    from telegram import __version_info__
except ImportError:
    __version_info__ = (0, 0, 0, 0, 0)  # type: ignore[assignment]

if __version_info__ < (20, 0, 0, "alpha", 1):
    raise RuntimeError(
        f"This example is not compatible with your current PTB version {TG_VER}. To view the "
        f"{TG_VER} version of this example, "
        f"visit https://docs.python-telegram-bot.org/en/v{TG_VER}/html"
    )

#connection = pika.BlockingConnection(pika.ConnectionParameters('localhost'))
# channel = connection.channel()
# channel.queue_declare(queue='generations')
# channel.queue_declare(queue='upscale')

def enqueue(connection, queue, task):
    return
    with connection.channel() as channel:
        channel.basic_publish(exchange='',
                        routing_key=queue,
                        body=pickle.dumps(task))

def callback_example(ch, method, properties, body):
    data = pickle.loads(body)
    print(" [x] Received %r" % data)


def consume(callback, queue):
    with connection.channel() as channel:
        channel.basic_consume(queue=queue,
            auto_ack=True,
            on_message_callback=callback)
        channel.start_consuming()


images_folder = "X:\\bot_generations"
airtable_token = getenv("AIRTABLE_TOKEN", "none")
api = Api(airtable_token)
# airtable_base = api.get_base(getenv("AIRTABLE_BASE_ID", "none"))
# generations = airtable_base.get_table("generations")
##nudenet_classifier = predict.load_model("./nsfw_mobilenet2/saved_model.h5")
session_id = str(uuid4())

# proxy = connection.ConnectionTcpMTProxyRandomizedIntermediate(
#     getenv("PROXY_HOST"), getenv("PROXY_PORT"), getenv("PROXY_SECRET"), session_id)



# Enable logging
logging.basicConfig(
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO
)
logger = logging.getLogger(__name__)
configs = dc.Cache("config")
black_list = dc.Cache("black_list")
generations_cache = dc.Cache("generations")
generations_files = dc.Cache("generations_files")
generations_meta = dc.Cache("generations_meta")

sentry_dsn = getenv("DSN", None)
# sentry_sdk.init(
#     sentry_dsn,
#     sample_rate=1,
#     traces_sample_rate=1.0,
#     profiles_sample_rate=1,
# )
sentry_sdk.capture_message("Service started", "info")


labels = {
    "prompt": "Затравка для генерации",
    "negative_prompt": "Негативные условия (чего быть не должно)",
    # "seed": "Зерно для генерации",
    # "width": "Ширина",
    # "height": "Высота",
    # "styles": None,
    # "cfg_scale": "Соответствие результата затравке",
    # "sampler_index": None,
    # "steps": "Глубина прорисовки",
    # "enable_hr": None,
    # "hr_scale": None,
    # "hr_upscaler": None,
    # "hr_second_pass_steps": "Увеличение детализации",
    # "hr_resize_x": None,
    # "hr_resize_y": None,
    # "denoising_strength": "Сила сглаживания перед детализацией"
}

# from PIL import Image, PngImagePlugin

# info = PngImagePlugin.PngInfo()
# info.add_text("text", "This is the text I stored in a png")
# info.add_text("ZIP", "VALUE", zip=True)
# im = Image.open("001.png")
# im.save("001.png", "PNG", pnginfo=info)
# im3 = Image.open("001.png")
# #print(im3.info["text"])
# print(im3.text["text"])

app_version = "1.0"

commands = []
menu_commands = []

for k, v in labels.items():
    if v:
        menu_commands.append(v)
        commands.append([InlineKeyboardButton(v, callback_data=k)])

reply_markup = InlineKeyboardMarkup(commands)

main_commands = [["/start"]]
reply_main_markup = ReplyKeyboardMarkup(main_commands, is_persistent=True)

# create API client with custom host, port
api = webuiapi.WebUIApi(host=getenv("API_HOST"), port=getenv("API_PORT"), use_https=True)
# all_options = api.get_options()


pf = ParquetFile('prompts.parquet')
df = pf.to_pandas()

def get_random_prompt():
    return df['prompt'].sample().tolist()[0]
    

async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    reply_markup = ReplyKeyboardRemove()
    init_params(update, context)
    await update.message.reply_text(
        """
    This chatbot is designed for a small group of friends who enjoy creating digital art using the stable diffusion neural generative model. However, any user can use it to create creative content. The bot is created to give an opportunity to those who want to create.
    But, please note that NSFW content is not welcome on this bot. There is a censorship filter here that monitors generation, and if it detects such content, it will issue a warning. After 5 warnings, the bot will block you. However, the filter may not always work correctly and if you received a warning on an image that is not related to NSFW, you can write to the bot @sd_bot_feedback_bot to get unblocked.
    We hope that you will create many beautiful and interesting images!
        """,
        reply_markup=reply_markup,
    )


def init_params(update, context):
    user = update.effective_user.name
    if user not in configs or configs.get(user, {}).get("version", None) != app_version:
        new_config = {}
        new_config["chat_id"] = context._chat_id
        new_config["version"] = app_version
        new_config["generation_params"] = default.generation_params
        configs[user] = new_config
        return configs[user]
    else:
        return configs[user]


async def send_admin(generation_id, update, user, generation_params, img_io, context):
    admin_chat_id = getenv("ADMIN_CHAT_ID", None)
    if generation_params:
        buttons = []
        generations_cache[generation_id] = generation_params
        generations_meta[generation_id] = {
            "username": user.name,
            "chat_id": context._chat_id,
            "last_gid": generation_id
        }
        buttons.append(
            [
                InlineKeyboardButton(
                    text="Re-generation", callback_data=f"regenerate:{generation_id}"
                )
            ]
        )
        # buttons.append(
        #     [
        #         InlineKeyboardButton(
        #             text="Upscale", callback_data=f"upscale:{generation_id}"
        #         )
        #     ]
        # )
        buttons.append(
            [
                InlineKeyboardButton(
                    text="Get prompt", callback_data=f"get_prompt:{generation_id}"
                )
            ]
        )
        buttons.append(
            [
                InlineKeyboardButton(
                    text="Strike", callback_data=f"strike:{generation_id}"
                )
            ]
        )
        keyboard = InlineKeyboardMarkup(buttons)
    else:
        keyboard = None
    if admin_chat_id and int(admin_chat_id) != context._chat_id:
        bot = update.get_bot()
        return await bot.send_photo(
            admin_chat_id,
            img_io,
            f"image from {user.name} {user.username} [#{generation_id}]",
            reply_markup=keyboard,
        )


STATE = None
if exists("blacklist.txt"):
    with open("blacklist.txt") as f:
        black_list_users = [x.strip() for x in f.readlines()]
else:
    black_list_users = []

if exists("whitelist.txt"):
    with open("whitelist.txt") as f:
        white_list_users = [x.strip() for x in f.readlines()]
else:
    white_list_users = []


async def echo(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    global api
    user = update.effective_user
    if user:
        init_params(update, context)  # Инициализация пользовательских настроек
        if not update.message or not update.message.text:
            return
        if 'paperspacegradient' in update.message.text:
            try:
                parsed_uri = urlparse(update.message.text)
                new_api_host = parsed_uri.hostname
                api = webuiapi.WebUIApi(host=new_api_host, port=getenv("API_PORT"), use_https=True)
                api.refresh_checkpoints()
                await update.message.reply_text(
                    "Настройки применены успешно",
                    reply_to_message_id=update.message.id,
                )
                return
            except Exception as e:
                await update.message.reply_text(
                    f"Не удалось применить настройки: {e}",
                    reply_to_message_id=update.message.id,
                )
        current_strike_count = black_list.get(user.name, 0)
        print(user.name, current_strike_count)
        if (
            user.name in black_list_users or current_strike_count >= 5
        ) and user.name not in white_list_users:
            print(user.name, "[blocked] user request rejected because user has banned")
            await update.message.reply_text(
                f"Strike {current_strike_count}/5 [Access to this bot is blocked for you for creating NSFW content.] If you think you were blocked by mistake, you can write to boot @sd_bot_feedback_out to get unblocked",
                reply_to_message_id=update.message.id,
            )
            return
        generation_params = copy(default.generation_params_low)
        translated = GoogleTranslator(source="auto", target="en").translate(
            update.message.text
        )
        if user.name not in white_list_users:
            censored_text = profanity.censor(translated)
        else:
            censored_text = translated
        generation_params["prompt"] = censored_text.replace("*", "")
        if " not " in translated.lower():
            generation_params["negative_prompt"] = ",".join(
                translated.lower().split(" not ")[1:]
            )
            generation_params["prompt"] = censored_text.lower().split(" not ")[0]
        if len(translated) < 3:
            return
        generation_params["seed"] = randint(1, 1000000)
        print(user.name, generation_params)
        configs[user.name] = generation_params
        try:
            await new_image(update, context, user, generation_params, censored_text)
        except Exception as e:
            await update.message.reply_text(f"An error has occurred: {e}")
            raise e


async def random(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    global api
    if ' ' in str(update.message.text):
        count = int(update.message.text.split(' ')[-1])
    else:
        count = 3
    user = update.effective_user
    if user:
        init_params(update, context)  # Инициализация пользовательских настроек
        for x in tqdm(range(count), desc=f"{user} N{count}"):
            generation_params = copy(default.generation_params_low)
            generation_params["prompt"] = get_random_prompt()
            generation_params["seed"] = randint(1, 1000000)
            print(user.name, generation_params)
            configs[user.name] = generation_params
            try:
                await new_image(update, context, user, generation_params, generation_params["prompt"])
            except Exception as e:
                await update.message.reply_text(f"An error has occurred: {e}")
                raise e

async def new_image(update, context, user, generation_params, censored_text, count=0):
    req_uid = str(uuid4())[:8]
    await user.send_chat_action(telegram.constants.ChatAction.TYPING)

    await user.send_message(
                "Generation request accepted", reply_to_message_id=update.message.id
            )
            
    img_io, filename = generate_image(
                user.name, generation_params, req_uid, generation_params["seed"]
            )

    await user.send_chat_action(telegram.constants.ChatAction.TYPING)
    await send_admin(
                req_uid, update, user, generation_params, img_io, context
            )
    img_io.seek(0)
    await user.send_chat_action(
                telegram.constants.ChatAction.UPLOAD_PHOTO
            )
    buttons = []
    generations_cache[req_uid] = generation_params
    buttons.append(
                [
                    InlineKeyboardButton(
                        text="Regenerate", callback_data=f"regenerate:{req_uid}"
                    )
                ]
            )
    buttons.append(
                [
                    InlineKeyboardButton(
                        text="Get prompt", callback_data=f"get_prompt:{req_uid}"
                    )
                ]
            )
    keyboard = InlineKeyboardMarkup(buttons)
    await update.message.reply_photo(
                img_io,
                f"[#{req_uid}] {censored_text}"[:min(len(censored_text), 300)],
                filename=f"{req_uid}.png",
                reply_to_message_id=update.message.id,
                reply_markup=keyboard,
            )
    remove(filename)


def check_filter(censored_text, censor_result, filter_name, filter_edge):
    if censor_result[filter_name] > filter_edge:
        block_reason = f"{filter_name}={round(censor_result[filter_name]*100)}%"
        censored_text = f"[BLOCKED {block_reason}] {censored_text}"
        return censored_text, True, round(censor_result[filter_name] * 100)
    else:
        return censored_text, False, round(censor_result[filter_name] * 100)


def generate_image(username, job_config, gen_id, seed):
    with sentry_sdk.start_transaction(
        op="task",
        name=f"[{username}]-{gen_id}-{job_config['seed']}",
        source="generate_image",
    ) as tran:
        with sentry_sdk.start_span(description="generate"):
            with tran.start_child(op="generate") as span:
                with sentry_sdk.start_span(description="txt2img"):
                    result1 = api.txt2img(**job_config)
                with sentry_sdk.start_span(description="save"):
                    img_io = BytesIO()
                    result1.image.save(img_io, "jpeg", optimize=False, quality=100)
                    img_io.seek(0)
                    filename = join(
                        images_folder, f"{username}-{gen_id}-{job_config['seed']}.jpg"
                    )
                    generations_files[f"{gen_id}-{job_config['seed']}"] = filename
                    with open(filename, "wb") as f:
                        result1.image.save(f, "jpeg", optimize=False, quality=100)
                    span.finish(tran.hub, end_timestamp=dt.now())
                    tran.set_status("success")
                    return img_io, filename


async def handle_callback(update, context):
    user = update.effective_user
    query = update.callback_query
    action, generation_id = query.data.split(":")
    generation_params = generations_cache[generation_id]
    meta = (
        generations_meta[generation_id] if generation_id in generations_meta else None
    )
    print(action.upper(), generation_params["prompt"])
    for _ in tqdm(
        range(5), desc=f"{datetime.now().isoformat()} Generate image for {user.name}"
    ):
        t = datetime.now()
        if action == "strike":
            strikes = black_list.get(user.name, 0) + 1
            black_list.set(user.name, strikes)
            message = f"Strike {strikes}/5 [Strike for NSFW content generation]"
            await update.callback_query.message.reply_text(message)
            bot = update.get_bot()
            await bot.send_message(
                meta["chat_id"],
                message,
            )
            return

        # meta = generations.create(
        #     {
        #         "username": user.name,
        #         "gid": generation_id,
        #         "prompt": generation_params["prompt"],
        #         "negative_prompt": generation_params["negative_prompt"],
        #         "seed": generation_params["seed"],
        #         "status": "processing",
        #         "hd": False,
        #     }
        # )
        if action == "max_upscale":
            generation_params.update(copy(default.generation_params_hq))
            await user.send_message(
                "Generation request accepted",
                reply_to_message_id=update.callback_query.message.id,
            )
            try:
                source_file = generations_files[
                    f"{generation_id}-{generation_params['seed']}"
                ]
                img_io, filename = await upscale_image(
                    generation_id, update, user, generation_params, source_file, context
                )
            except Exception as e:
                sentry_sdk.capture_exception(e)
                await user.send_message(
                    "Generation failed: the server is under maintenance, try to send a request later.",
                    reply_to_message_id=update.callback_query.message.id,
                )
                return
            await update.callback_query.message.reply_photo(
                img_io,
                f"MAX-Upscale: #{generation_id}, seed {generation_params['seed']}",
                filename=f"{filename}.png",
                reply_to_message_id=update.callback_query.message.id,
            )
            delta = (datetime.now() - t).total_seconds()
            # generations.update(
            #     meta["id"],
            #     {"status": "done", "duration": delta, "filename": filename, "hd": True},
            # )
        if action == "regenerate":
            generation_params["seed"] = randint(1, 1000000)
            await user.send_message(
                "Generation request accepted",
                reply_to_message_id=update.callback_query.message.id,
            )
            try:
                img_io, filename = await create_new_image(
                    generation_id, update, user, generation_params, context
                )
            except Exception as e:
                sentry_sdk.capture_exception(e)
                await user.send_message(
                    "Generation failed: the server is under maintenance, try to send a request later.",
                    reply_to_message_id=update.callback_query.message.id,
                )
                return
            buttons = []
            generations_cache[generation_id] = generation_params
            buttons.append(
                [
                    InlineKeyboardButton(
                        text="Re-generation",
                        callback_data=f"regenerate:{generation_id}",
                    )
                ]
            )
            buttons.append(
                [
                    InlineKeyboardButton(
                        text="Get prompt", callback_data=f"get_prompt:{generation_id}"
                    )
                ]
            )
            # if not generation_params["enable_hr"]:
            #     buttons.append(
            #         [
            #             InlineKeyboardButton(
            #                 text="Upscale", callback_data=f"upscale:{generation_id}"
            #             )
            #         ]
            #     )
            # else:
            #     buttons.append(
            #         [
            #             InlineKeyboardButton(
            #                 text="Max upscale",
            #                 callback_data=f"max_upscale:{generation_id}",
            #             )
            #         ]
            #     )
            keyboard = InlineKeyboardMarkup(buttons)
            await update.callback_query.message.reply_photo(
                img_io,
                f"Re-generation: #{generation_id}, seed {generation_params['seed']}",
                filename=f"{filename}.png",
                reply_to_message_id=update.callback_query.message.id,
                reply_markup=keyboard,
            )
            delta = (datetime.now() - t).total_seconds()
            # generations.update(
            #     meta["id"],
            #     {
            #         "status": "done",
            #         "duration": delta,
            #         "filename": filename,
            #         "seed": generation_params["seed"],
            #         "hd": True,
            #     },
            # )
        if action == "get_prompt":
            negative_prompt = (
                f" not {generation_params['negative_prompt']}"
                if len(generation_params["negative_prompt"]) > 1
                else ""
            )
            await user.send_message(
                f"{generation_params['prompt']}{negative_prompt}",
                reply_to_message_id=update.callback_query.message.id,
            )
            return


async def create_new_image(generation_id, update, user, generation_params, context):
    await user.send_chat_action(telegram.constants.ChatAction.TYPING)
    img_io, filename = generate_image(
        user.name, generation_params, generation_id, generation_params["seed"]
    )
    await user.send_chat_action(telegram.constants.ChatAction.UPLOAD_PHOTO)
    img_io, filename = upscale_image_extra(
        user.name, generation_params, generation_id, filename
    )
    await send_admin(generation_id, update, user, generation_params, img_io, context)
    img_io.seek(0)
    return img_io, filename


async def upscale_image(
    generation_id, update, user, generation_params, source_file, context
):
    await user.send_chat_action(telegram.constants.ChatAction.TYPING)
    img_io, filename = upscale_image_extra(
        user.name, generation_params, generation_id, source_file
    )
    await user.send_chat_action(telegram.constants.ChatAction.TYPING)
    await send_admin(generation_id, update, user, generation_params, img_io, context)
    img_io.seek(0)
    await user.send_chat_action(telegram.constants.ChatAction.UPLOAD_PHOTO)
    return img_io, filename


def upscale_image_extra(username, job_config, gen_id, source_file):
    with sentry_sdk.start_transaction(
        op="task",
        name=f"[{username}]-{gen_id}-{job_config['seed']}",
        source="generate_image",
    ) as tran:
        with sentry_sdk.start_span(description="generate"):
            with tran.start_child(op="generate") as span:
                with sentry_sdk.start_span(description="txt2img"):
                    with Image.open(source_file) as img:
                        result1 = api.extra_single_image(img, upscaler_1="Lanczos", upscaling_resize=4)  # type: ignore
                with sentry_sdk.start_span(description="save"):
                    img_io = BytesIO()
                    result1.image.save(img_io, "jpeg", optimize=False, quality=98)
                    img_io.seek(0)
                    filename = join(
                        images_folder, f"{username}-{gen_id}-{job_config['seed']}.png"
                    )
                    generations_files[f"{gen_id}-{job_config['seed']}"] = filename
                    with open(filename, "wb") as f:
                        result1.image.save(f, "jpeg", optimize=False, quality=98)
                    span.finish(tran.hub, end_timestamp=dt.now())
                    tran.set_status("success")
                    return img_io, filename


async def photo_to_art(update, context):   
    user = update.effective_user
    #generation_params = configs[user.name]
    generation = copy(default.generation_params_img2img)
    #generation['prompt'] = generation_params['prompt']
    #generation['negative_prompt'] = generation_params['negative_prompt']
    generation['seed'] = randint(0, 100000)
    fileID = update.message.photo[-1].file_id 
    bot = update.get_bot()  
    file_info = await bot.get_file(fileID)
    source_file = f"img2img_{fileID}.png"
    await file_info.download_to_drive(source_file)
    for x in tqdm(range(1), desc=f"img2img for {user.name}"):
        with Image.open(source_file) as img:
            #crop = img.thumbnail((512, 512), Image.LANCZOS)
            interrogate_result = api.interrogate(img)
            image_info = interrogate_result.info
            generation['prompt'] = image_info
            result1 = api.img2img(images=[img], mask_image=None, **generation, restore_faces=True)
            fixed_width = generation['width']
            width_percent = (fixed_width / float(img.width))
            height_size = int((float(img.size[0]) * float(width_percent)))
            generation['width'] = fixed_width
            generation['height'] = height_size
            img_io = BytesIO()
            result1.image.save(img_io, "jpeg", optimize=False, quality=100)
            img_io.seek(0)
            filename = join(
                images_folder, f"{user.name}-{fileID}.png"
            )
            generations_files[f"{fileID}"] = filename
            with open(filename, "wb") as f:
                result1.image.save(f, "jpeg", optimize=False, quality=100)
            with Image.open(filename) as img:
                result1 = api.extra_single_image(img, upscaler_1="4x-UltraSharp", upscaling_resize=2, upscale_first=True)  # type: ignore
                img_io = BytesIO()
                result1.image.save(img_io, "jpeg", optimize=False, quality=100)
                img_io.seek(0)
            await user.send_chat_action(telegram.constants.ChatAction.UPLOAD_PHOTO)
            await update.message.reply_photo(
                    img_io,
                    f"Img2img: #{fileID}, seed {generation['seed']}: {image_info}",
                    filename=f"{filename}.jpg",
                    reply_to_message_id=update.message.id
                )
            
        

def upscale(ch, method, properties, body):
    data = pickle.loads(body)
    print(" [x] Received %r" % data)
    # await user.send_message(
    #     "Generation request accepted",
    #     reply_to_message_id=update.callback_query.message.id,
    # )
    # try:
    #     img_io, filename = await create_new_image(
    #         generation_id, update, user, generation_params, context
    #     )
    # except Exception as e:
    #     sentry_sdk.capture_exception(e)
    #     await user.send_message(
    #         "Generation failed: the server is under maintenance, try to send a request later.",
    #         reply_to_message_id=update.callback_query.message.id,
    #     )
    #     return
    # buttons = []
    # generations_cache[generation_id] = generation_params
    # buttons.append(
    #     [
    #         InlineKeyboardButton(
    #             text="Re-generation",
    #             callback_data=f"regenerate:{generation_id}",
    #         )
    #     ]
    # )
    # # buttons.append(
    # #     [
    # #         InlineKeyboardButton(
    # #             text="Max upscale", callback_data=f"max_upscale:{generation_id}"
    # #         )
    # #     ]
    # # )
    # buttons.append(
    #     [
    #         InlineKeyboardButton(
    #             text="Get prompt", callback_data=f"get_prompt:{generation_id}"
    #         )
    #     ]
    # )
    # keyboard = InlineKeyboardMarkup(buttons)
    # await update.callback_query.message.reply_photo(
    #     img_io,
    #     f"Upscale: #{generation_id}, seed {generation_params['seed']}",
    #     filename=f"{filename}.png",
    #     reply_to_message_id=update.callback_query.message.id,
    #     reply_markup=keyboard,
    # )
    # delta = (datetime.now() - t).total_seconds()
    # generations.update(
    #     meta["id"],
    #     {"status": "done", "duration": delta, "filename": filename, "hd": True},
    # )

bot = Application.builder().token(getenv("TOKEN")).build()  # type: ignore

def main() -> None:
    """Start the bot."""
    # Create the Application and pass it your bot's token.
    # thread = Thread(target = consume, args = (upscale, "upscale"))
    # thread.start()


    # on different commands - answer in Telegram
    bot.add_handler(CommandHandler("start", start))
    bot.add_handler(CommandHandler("random", random))

    # on non command i.e message - echo the message on Telegram
    bot.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, echo))

    bot.add_handler(MessageHandler(filters.PHOTO, photo_to_art))
    # application.add_handler(InlineQueryHandler(echo))
    # application.add_handler(MessageHandler(filters.COMMAND, echo))
    bot.add_handler(CallbackQueryHandler(handle_callback))
    # Run the bot until the user presses Ctrl-C
    while True:
        try:
            bot.run_polling(poll_interval=3)
        except Exception as e:
            print(e)


if __name__ == "__main__":
    main()