diff --git a/main.py b/main.py index dd6a607..c576f6b 100644 --- a/main.py +++ b/main.py @@ -2,7 +2,7 @@ import json import requests import io import os -import random +import uuid import base64 from PIL import Image, PngImagePlugin from pyrogram import Client, filters @@ -26,55 +26,81 @@ SD_URL = os.environ.get("SD_URL", None) print(SD_URL) app = Client("stable", api_id=API_ID, api_hash=API_HASH, bot_token=TOKEN) +#default params +steps_value_default = 40 -def slice_positive_negative(string): - delimiter = "ng:" - if delimiter in string: - index = string.index(delimiter) - positive = string[:index].rstrip() - negative = string[index + len(delimiter) :].lstrip() - return positive, negative +def process_input_string(string): + ng_delimiter = "ng:" + steps_delimiter = "steps:" + + ng_index = string.find(ng_delimiter) + steps_index = string.find(steps_delimiter) + + if ng_index != -1 and steps_index != -1: + if ng_index < steps_index: + positive = string[:ng_index].strip() + negative = string[ng_index + len(ng_delimiter):steps_index].strip() + steps_str = string[steps_index + len(steps_delimiter):].strip().split()[0] + else: + positive = string[:steps_index].strip() + negative = string[steps_index + len(steps_delimiter):ng_index].strip() + steps_str = string[ng_index + len(ng_delimiter):].strip().split()[0] + elif ng_index != -1: + positive = string[:ng_index].strip() + negative = string[ng_index + len(ng_delimiter):].strip() + steps_str = None + elif steps_index != -1: + positive = string[:steps_index].strip() + negative = None + steps_str = string[steps_index + len(steps_delimiter):].strip().split()[0] else: - return string, "null" # if we are missing the ng bad thing happen.... - #POSSIBLE NG [[[[[[[[[[lowres, low resolution, bad resolution, bad, poor quality, bad quality, blurry, bad art, incomplete, logo, signature, text, jpeg artifacts, compression artifacts,child, childish]]]]]]]]]] + positive = string.strip() + negative = None + steps_str = None + + try: + steps_value = int(steps_str) + #limit steps to range + if not 1 <= steps_value <= 70: + steps_value = steps_value_default + except (ValueError, TypeError): + steps_value = None + + + return positive, negative, steps_value @app.on_message(filters.command(["draw"])) def draw(client, message): msgs = message.text.split(" ", 1) if len(msgs) == 1: message.reply_text( - "Format :\n/draw < text to image >\nng: < negative (optional) >" + "Format :\n/draw < text to image >\nng: < negative (optional) >\nsteps: < steps value (1-70, optional) >" ) return - msg = slice_positive_negative(msgs[1]) - print(msg) + positive, negative, steps_value = process_input_string(msgs[1]) payload = { - "prompt": msg[0], - "negative_prompt": msg[1], + "prompt": positive, } + if negative is not None: + payload["negative_prompt"] = negative + if steps_value is not None: + payload["steps"] = steps_value print(payload) + # The rest of the draw function remains unchanged + + K = message.reply_text("Please Wait 10-15 Second") r = requests.post(url=f"{SD_URL}/sdapi/v1/txt2img", json=payload).json() def genr(): - chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" - chars1 = "1234564890" - gen1 = random.choice(chars) - gen2 = random.choice(chars) - gen3 = random.choice(chars1) - gen4 = random.choice(chars) - gen5 = random.choice(chars) - gen6 = random.choice(chars) - gen7 = random.choice(chars1) - gen8 = random.choice(chars) - gen9 = random.choice(chars) - gen10 = random.choice(chars1) - return f"{message.from_user.id}-MOE{gen1}{gen2}{gen3}{gen4}{gen5}{gen6}{gen7}{gen8}{gen9}{gen10}" + unique_id = str(uuid.uuid4())[:7] + return f"{message.from_user.first_name}-{unique_id}" word = genr() + for i in r["images"]: image = Image.open(io.BytesIO(base64.b64decode(i.split(",", 1)[0]))) @@ -86,9 +112,15 @@ def draw(client, message): image.save(f"{word}.png", pnginfo=pnginfo) message.reply_photo( - photo=f"{word}.png", - caption=f"Prompt - **{msg[0]}**\nNegative Prompt - **{msg[1]}**\n**[{message.from_user.first_name}-Kun](tg://user?id={message.from_user.id})**", - ) + photo=f"{word}.png", + caption=( + f"Prompt - **{positive}**\n" + f"Negative Prompt - **{negative if negative is not None else 'None'}**\n" + f"Steps - **{steps_value if steps_value != steps_value_default else 'Default'}**\n" + f"**[{message.from_user.first_name}-Kun](tg://user?id={message.from_user.id})**" + ), +) + # os.remove(f"{word}.png") K.delete()