add steps and nicer filename

This commit is contained in:
yair 2023-04-28 10:02:47 +03:00
parent 9bce6f989e
commit 095e418da9

94
main.py
View File

@ -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()