# Copyright 2023 Qewertyy, MIT License
import asyncio
from lexica import AsyncClient

async def ImageGeneration(model,prompt):
    try:
        client = AsyncClient()
        output = await client.generate(model,prompt,"")
        if output['code'] != 1:
            return 2
        elif output['code'] == 69:
            return output['code']
        task_id, request_id = output['task_id'],output['request_id']
        await asyncio.sleep(20)
        tries = 0
        image_url = None
        resp = await client.getImages(task_id,request_id)
        while True:
            if resp['code'] == 2:
                image_url = resp['img_urls']
                break
            if tries > 15:
                break
            await asyncio.sleep(5)
            resp = await client.getImages(task_id,request_id)
            tries += 1
            continue
        return image_url
    except Exception as e:
        raise Exception(f"Failed to generate the image: {e}")
    finally:
        await client.close()

async def UpscaleImages(image: bytes) -> str:
    """
    Upscales an image and return with upscaled image path.
    """
    try:
        client = AsyncClient()
        content = await client.upscale(image)
        await client.close()
        upscaled_file_path = "upscaled.png"
        with open(upscaled_file_path, "wb") as output_file:
            output_file.write(content)
        return upscaled_file_path
    except Exception as e:
        raise Exception(f"Failed to upscale the image: {e}")

async def gpt(prompt) -> str:
    try:
        client = AsyncClient()
        output = await client.gpt(prompt)
        await client.close()
        return output['content']
    except Exception as E:
        raise Exception(f"Failed to upscale the image: {e}")
Utils/API.py
Qewertyy/SDWaifuRobot
from .API import *
from .btnParser import *
from .miscs import *
from .pastebins import *
Utils/__init__.py
Qewertyy/SDWaifuRobot
from pyrogram.types import InlineKeyboardButton
from math import ceil

class EqInlineKeyboardButton(InlineKeyboardButton):
    def __eq__(self, other):
        return self.text == other.text

    def __lt__(self, other):
        return self.text < other.text

    def __gt__(self, other):
        return self.text > other.text

def paginate_models(page_n: int, models: list,user_id) -> list:
    modules = sorted(
        [
            EqInlineKeyboardButton(
            x['name'],
            callback_data=f"d.{x['id']}.{user_id}"
                )
                for x in models
            ]
            )

    pairs = list(zip(modules[::3], modules[1::3]))
    i = 0
    for m in pairs:
        for _ in m:
            i += 1
    if len(modules) - i == 1:
        pairs.append((modules[-1],))
    elif len(modules) - i == 2:
        pairs.append(
            (
                modules[-2],
                modules[-1],
            )
        )

    COLUMN_SIZE = 3

    max_num_pages = ceil(len(pairs) / COLUMN_SIZE)
    modulo_page = page_n % max_num_pages

    # can only have a certain amount of buttons side by side
    if len(pairs) > COLUMN_SIZE:
        pairs = pairs[
            modulo_page * COLUMN_SIZE : COLUMN_SIZE * (modulo_page + 1)
        ] + [
            (
                EqInlineKeyboardButton(
                    "⬅️",
                    callback_data=f"d.left.{modulo_page}.{user_id}"
                ),
                EqInlineKeyboardButton(
                    "Cancel",
                    callback_data=f"d.-1.{user_id}"
                ),
                EqInlineKeyboardButton(
                    "➡️",
                    callback_data=f"d.right.{modulo_page}.{user_id}"
                ),
            )
        ]
    else:
        pairs += [[EqInlineKeyboardButton("Cancel", callback_data=f"d.-1.{user_id}")]]

    return pairs
Utils/btnParser.py
Qewertyy/SDWaifuRobot
from functools import wraps

def errorHandler(func):
    @wraps(func)
    async def wrapper(client, message, *args, **kwargs):
        try:
            await func(client, message, *args, **kwargs)
        except Exception as err:
            exc_type, exc_obj, exc_tb = sys.exc_info()
            errors = traceback.format_exception(
                exc_type,
                exc_obj,
                exc_tb,
            )
            errors = evaluateContent(
                "#ERROR | `{}` | `{}`\n\n`{}`\n```{}```\n".format(
                    0 if not message.from_user else message.from_user.id,
                    0 if not message.chat else message.chat.id,
                    message.caption if message.caption else message.text,
                    "".join(errors),
                ),
            )
            await client.send_message(Config.LOG_CHANNEL, errors)
            raise err
    return wrapper
Utils/decorators.py
Qewertyy/SDWaifuRobot
# Copyright 2023 Qewertyy, MIT License

async def getFile(message):
    if not message.reply_to_message:
        return None
    if message.reply_to_message.document is False or message.reply_to_message.photo is False:
        return None
    if message.reply_to_message.document and message.reply_to_message.document.mime_type in ['image/png','image/jpg','image/jpeg'] or message.reply_to_message.photo:
        image = await message.reply_to_message.download()
        return image
    else:
        return None

def getText(message):
    """Extract Text From Commands"""
    text_to_return = message.text
    if message.text is None:
        return None
    if " " in text_to_return:
        try:
            return message.text.split(None, 1)[1]
        except IndexError:
            return None
    else:
        return None
Utils/miscs.py
Qewertyy/SDWaifuRobot
# Copyright 2023 Qewertyy, MIT License
import httpx

async def nekobin(data):
    """
    To Paste the given message/text/code to nekobin
    """
    try:
        async with httpx.AsyncClient() as req:
            res = req.post(
                url=NEKOBIN,
                json={
                    "content":data,
                    "title": "data",
                    "author": "SDWaifuRobot"
                })
    except Exception as e:
        return {"error": str(e)}
    if res.ok:
        resp = res.json()
        purl = (
            f"nekobin.com/{resp['result']['key']}.{extension}"
            if extension
            else f"nekobin.com/{resp['result']['key']}"
        )
        return purl
    return {"error": "Unable to reach nekobin."}
Utils/pastebins.py
Qewertyy/SDWaifuRobot
# Copyright 2023 Qewertyy, MIT License

import uvloop
uvloop.install()
import datetime,logging, sys
from pyrogram import Client
from lexica import Client as ApiClient
from config import Config

# Get logging configurations
logging.basicConfig(
    format="%(asctime)s - [BOT] - %(levelname)s - %(message)s",
    handlers=[logging.FileHandler("logs.txt"), logging.StreamHandler()],
    level=logging.INFO,
)
logging.getLogger("pyrogram").setLevel(logging.WARNING)
logging.getLogger("httpx").setLevel(logging.WARNING)
LOGGER = logging.getLogger(__name__)

StartTime = Models = None

class Bot(Client):
    global StartTime,Models
    StartTime = datetime.datetime.now()
    api = ApiClient()
    Models = api.getModels()['models']['image']
    def __init__(self):
        super().__init__(
            "SDWaifuRobot",
            api_id=Config.API_ID,
            api_hash=Config.API_HASH,
            bot_token=Config.BOT_TOKEN,
            plugins=dict(root="plugins"),
        )
    async def start(self):
        await super().start()
        LOGGER.info("Bot Started")

    if Models is None:
        LOGGER.error("Models are empty!")
        sys.exit(1)

    async def stop(self):
        await super().stop()
        LOGGER.info("Stopped Services")

if __name__ == "__main__":
    Bot().run()
bot.py
Qewertyy/SDWaifuRobot
# Copyright 2023 Qewertyy, MIT License

from pyrogram import Client, filters, types as t
from Utils import getText,paginate_models,ImageGeneration
from bot import Models
Database = {}

@Client.on_message(filters.command(["draw","create","imagine","dream"]))
async def draw(_: Client, m: t.Message):
    global Database
    prompt = getText(m)
    if prompt is None:
        return await m.reply_text("give something to create")
    user = m.from_user
    data = {'prompt':prompt,'reply_to_id':m.id}
    Database[user.id] = data
    btns = paginate_models(0,Models,user.id)
    await m.reply_text(
            text=f"Your prompt: `{prompt}`\n\nSelect a model",
            reply_markup=t.InlineKeyboardMarkup(btns)
            )

@Client.on_callback_query(filters.regex(pattern=r"^d.(.*)"))
async def selectModel(_:Client,query:t.CallbackQuery):
    global Database
    data = query.data.split('.')
    auth_user = int(data[-1])
    if query.from_user.id != auth_user:
        return await query.answer("No.")
    if len(data) > 3:
        if data[1] == "right":
            next_page = int(data[2])
            await query.edit_message_reply_markup(
                t.InlineKeyboardMarkup(
                    paginate_models(next_page + 1,Models,auth_user)
                    )
                )
        elif data[1] == "left":
            curr_page = int(data[2])
            await query.edit_message_reply_markup(
                t.InlineKeyboardMarkup(
                    paginate_models(curr_page - 1,Models,auth_user)
                )
            )
        return
    modelId = int(data[1])
    if modelId == -1:
        del Database[auth_user]
        await query.message.delete()
        return
    await query.edit_message_text("Please wait, generating your image")
    promptData = Database.get(auth_user,None)
    if promptData is None:
        return await query.edit_message_text("Something went wrong.")
    img_url = await ImageGeneration(modelId,promptData['prompt'])
    if img_url is None or img_url == 2 or img_url ==1:
        return await query.edit_message_text("something went wrong!")
    elif img_url == 69:
        return await query.edit_message_text("NSFW not allowed!")
    images = []
    modelName = [i['name'] for i in Models if i['id'] == modelId]
    for i in img_url:
        images.append(t.InputMediaDocument(i))
    images[-1] = t.InputMediaDocument(img_url[-1],caption=f"Your prompt: `{promptData['prompt']}`\nModel: `{modelName}`") # for caption
    await query.message.delete()
    try:
        del Database[auth_user]
    except KeyError:
        pass
    await _.send_media_group(
        chat_id=query.message.chat.id,
        media=images,
        reply_to_message_id=promptData['reply_to_id']
    )
plugins/draw.py
Qewertyy/SDWaifuRobot
# Copyright 2023 Qewertyy, MIT License

from pyrogram import Client, filters, types as t
from bot import StartTime

@Client.on_message(filters.command(["start","help","repo"]))
async def start(_: Client, m: t.Message):
    await m.reply_text(
        "Just an image generation bot by @Qewertyy.\n\nCommands: \n`/draw`: create images\n`/upscale`: upscale your images",
        reply_markup=t.InlineKeyboardMarkup(
            [
                [
                    t.InlineKeyboardButton(text="Source",url="https://github.com/Qewertyy/SDWaifuRobot")
                ]
            ]
        )
    )

plugins/start.py
Qewertyy/SDWaifuRobot
from pyrogram import Client, filters, types as t
from Utils import getFile, UpscaleImages
import os

@Client.on_message(filters.command(["upscale"]))
async def upscaleImages(_, message):
    file = await getFile(message)
    if file is None:
        return await message.reply_text("Replay to an image?")
    msg = await message.reply("wait a min...")
    imageBytes = open(file,"rb").read()
    os.remove(file)
    upscaledImage = await UpscaleImages(imageBytes)
    try:
      await message.reply_document(open(upscaledImage,"rb"))
      await msg.delete()
      os.remove(upscaledImage)
    except Exception as e:
       await msg.edit(f"{e}")
plugins/upscale.py
Qewertyy/SDWaifuRobot