2023-01-12 06:37:56 +02:00
import json
import requests
import io
2023-05-06 18:59:02 +03:00
import re
2023-03-18 00:04:54 +02:00
import os
2023-04-28 10:02:47 +03:00
import uuid
2023-01-12 06:37:56 +02:00
import base64
2023-03-18 00:04:54 +02:00
from PIL import Image , PngImagePlugin
from pyrogram import Client , filters
from pyrogram . types import *
2023-03-18 00:24:46 +02:00
from dotenv import load_dotenv
2023-03-18 00:04:54 +02:00
2023-04-22 22:18:10 +03:00
# Done! Congratulations on your new bot. You will find it at
2023-03-18 00:04:54 +02:00
# t.me/gootmornbot
# You can now add a description, about section and profile picture for your bot, see /help for a list of commands. By the way, when you've finished creating your cool bot, ping our Bot Support if you want a better username for it. Just make sure the bot is fully operational before you do this.
# Use this token to access the HTTP API:
# Keep your token secure and store it safely, it can be used by anyone to control your bot.
# For a description of the Bot API, see this page: https://core.telegram.org/bots/api
2023-01-12 06:37:56 +02:00
2023-03-18 00:24:46 +02:00
load_dotenv ( )
2023-04-22 22:18:10 +03:00
API_ID = os . environ . get ( " API_ID " , None )
API_HASH = os . environ . get ( " API_HASH " , None )
TOKEN = os . environ . get ( " TOKEN " , None )
SD_URL = os . environ . get ( " SD_URL " , None )
2023-03-18 00:04:54 +02:00
print ( SD_URL )
2023-04-22 22:18:10 +03:00
app = Client ( " stable " , api_id = API_ID , api_hash = API_HASH , bot_token = TOKEN )
2023-01-12 06:37:56 +02:00
2023-05-06 18:59:02 +03:00
# default params
2023-04-28 10:02:47 +03:00
steps_value_default = 40
2023-05-06 18:59:02 +03:00
def parse_input ( input_string ) :
default_payload = {
" prompt " : " " ,
" negative_prompt " : " " ,
" controlnet_input_image " : [ ] ,
" controlnet_mask " : [ ] ,
" controlnet_module " : " " ,
" controlnet_model " : " " ,
" controlnet_weight " : 1 ,
" controlnet_resize_mode " : " Scale to Fit (Inner Fit) " ,
" controlnet_lowvram " : False ,
" controlnet_processor_res " : 64 ,
" controlnet_threshold_a " : 64 ,
" controlnet_threshold_b " : 64 ,
" controlnet_guidance " : 1 ,
" controlnet_guessmode " : True ,
" enable_hr " : False ,
" denoising_strength " : 0.5 ,
" hr_scale " : 1.5 ,
" hr_upscale " : " Latent " ,
" seed " : - 1 ,
" subseed " : - 1 ,
" subseed_strength " : - 1 ,
" sampler_index " : " " ,
" batch_size " : 1 ,
" n_iter " : 1 ,
" steps " : 20 ,
" cfg_scale " : 7 ,
" width " : 512 ,
" height " : 512 ,
" restore_faces " : True ,
" override_settings " : { } ,
" override_settings_restore_afterwards " : True ,
}
# Initialize an empty payload with the 'prompt' key
payload = { " prompt " : " " }
prompt = [ ]
# Find all occurrences of keys (words ending with a colon)
matches = re . finditer ( r " ( \ w+): " , input_string )
last_index = 0
# Iterate over the found keys
for match in matches :
2023-05-06 21:18:30 +03:00
key = match . group ( 1 ) . lower ( ) # Convert the key to lowercase
2023-05-06 18:59:02 +03:00
value_start_index = match . end ( )
# If there's text between the last key and the current key, add it to the prompt
if last_index != match . start ( ) :
prompt . append ( input_string [ last_index : match . start ( ) ] . strip ( ) )
last_index = value_start_index
# Check if the key is in the default payload
if key in default_payload :
# Extract the value for the current key
value_end_index = re . search (
r " (?= \ s+ \ w+:|$) " , input_string [ value_start_index : ]
) . start ( )
value = input_string [
value_start_index : value_start_index + value_end_index
] . strip ( )
# Check if the default value for the key is an integer
if isinstance ( default_payload [ key ] , int ) :
# If the value is a valid integer, store it as an integer in the payload
if value . isdigit ( ) :
payload [ key ] = int ( value )
else :
# If the default value for the key is not an integer, store the value as is in the payload
payload [ key ] = value
last_index + = value_end_index
2023-04-28 10:02:47 +03:00
else :
2023-05-06 18:59:02 +03:00
# If the key is not in the default payload, add it to the prompt
prompt . append ( f " { key } : " )
# Join the prompt words and store it in the payload
payload [ " prompt " ] = " " . join ( prompt )
2023-05-06 21:09:46 +03:00
# If the prompt is empty, set the input string as the prompt
2023-05-06 18:59:02 +03:00
if not payload [ " prompt " ] :
2023-05-06 21:09:46 +03:00
payload [ " prompt " ] = input_string . strip ( )
2023-05-06 18:59:02 +03:00
# Return the final payload
return payload
2023-03-18 00:04:54 +02:00
2023-01-12 06:37:56 +02:00
@app.on_message ( filters . command ( [ " draw " ] ) )
def draw ( client , message ) :
2023-04-22 22:18:10 +03:00
msgs = message . text . split ( " " , 1 )
2023-01-12 06:37:56 +02:00
if len ( msgs ) == 1 :
2023-04-22 22:18:10 +03:00
message . reply_text (
2023-04-28 10:02:47 +03:00
" Format : \n /draw < text to image > \n ng: < negative (optional) > \n steps: < steps value (1-70, optional) > "
2023-04-22 22:18:10 +03:00
)
2023-01-12 06:37:56 +02:00
return
2023-05-06 18:59:02 +03:00
payload = parse_input ( msgs [ 1 ] )
2023-03-18 12:58:05 +02:00
print ( payload )
2023-01-12 06:37:56 +02:00
2023-04-28 10:02:47 +03:00
# The rest of the draw function remains unchanged
2023-04-22 22:15:16 +03:00
K = message . reply_text ( " Please Wait 10-15 Second " )
2023-04-22 22:18:10 +03:00
r = requests . post ( url = f " { SD_URL } /sdapi/v1/txt2img " , json = payload ) . json ( )
2023-01-12 06:37:56 +02:00
2023-03-18 00:24:46 +02:00
def genr ( ) :
2023-04-28 10:02:47 +03:00
unique_id = str ( uuid . uuid4 ( ) ) [ : 7 ]
return f " { message . from_user . first_name } - { unique_id } "
2023-04-22 22:18:10 +03:00
2023-03-18 12:58:05 +02:00
word = genr ( )
2023-04-28 10:02:47 +03:00
2023-04-22 22:18:10 +03:00
for i in r [ " images " ] :
2023-01-12 06:37:56 +02:00
image = Image . open ( io . BytesIO ( base64 . b64decode ( i . split ( " , " , 1 ) [ 0 ] ) ) )
png_payload = { " image " : " data:image/png;base64, " + i }
2023-04-22 22:18:10 +03:00
response2 = requests . post ( url = f " { SD_URL } /sdapi/v1/png-info " , json = png_payload )
2023-01-12 06:37:56 +02:00
pnginfo = PngImagePlugin . PngInfo ( )
pnginfo . add_text ( " parameters " , response2 . json ( ) . get ( " info " ) )
2023-04-22 22:18:10 +03:00
image . save ( f " { word } .png " , pnginfo = pnginfo )
2023-01-12 06:37:56 +02:00
2023-05-06 21:09:46 +03:00
# Add a flag to check if the user provided a seed value
user_provided_seed = " seed " in payload
2023-05-06 18:59:02 +03:00
info_dict = response2 . json ( )
seed_value = info_dict [ ' info ' ] . split ( " , Seed: " ) [ 1 ] . split ( " , " ) [ 0 ]
# print(seed_value)
caption = f " **[ { message . from_user . first_name } -Kun](tg://user?id= { message . from_user . id } )** \n \n "
for key , value in payload . items ( ) :
caption + = f " { key . capitalize ( ) } - ** { value } ** \n "
caption + = f " Seed - ** { seed_value } ** \n "
2023-01-12 06:37:56 +02:00
message . reply_photo (
2023-05-06 18:59:02 +03:00
photo = f " { word } .png " ,
caption = caption ,
)
2023-04-28 10:02:47 +03:00
2023-03-18 12:58:05 +02:00
# os.remove(f"{word}.png")
2023-01-12 06:37:56 +02:00
K . delete ( )
2023-03-18 00:04:54 +02:00
@app.on_message ( filters . command ( [ " getmodels " ] ) )
async def get_models ( client , message ) :
2023-04-22 22:18:10 +03:00
response = requests . get ( url = f " { SD_URL } /sdapi/v1/sd-models " )
2023-03-18 00:04:54 +02:00
if response . status_code == 200 :
models_json = response . json ( )
# create buttons for each model name
buttons = [ ]
for model in models_json :
2023-04-22 22:18:10 +03:00
buttons . append (
[
InlineKeyboardButton (
model [ " title " ] , callback_data = model [ " model_name " ]
)
]
)
2023-03-18 00:04:54 +02:00
# send the message
await message . reply_text (
text = " Select a model [checkpoint] to use " ,
2023-04-22 22:18:10 +03:00
reply_markup = InlineKeyboardMarkup ( buttons ) ,
2023-03-18 00:04:54 +02:00
)
2023-04-22 22:18:10 +03:00
2023-03-18 00:04:54 +02:00
@app.on_callback_query ( )
async def process_callback ( client , callback_query ) :
# if a model button is clicked, set sd_model_checkpoint to the selected model's title
sd_model_checkpoint = callback_query . data
# The sd_model_checkpoint needs to be set to the title from /sdapi/v1/sd-models
# post using /sdapi/v1/options
2023-04-22 22:18:10 +03:00
options = { " sd_model_checkpoint " : sd_model_checkpoint }
2023-03-18 00:04:54 +02:00
# post the options
2023-04-22 22:18:10 +03:00
response = requests . post ( url = f " { SD_URL } /sdapi/v1/options " , json = options )
2023-03-18 00:04:54 +02:00
if response . status_code == 200 :
# if the post was successful, send a message
2023-04-22 22:18:10 +03:00
await callback_query . message . reply_text (
" checpoint set to " + sd_model_checkpoint
)
2023-03-18 00:04:54 +02:00
else :
# if the post was unsuccessful, send an error message
await callback_query . message . reply_text ( " Error setting options " )
2023-01-12 06:37:56 +02:00
@app.on_message ( filters . command ( [ " start " ] , prefixes = [ " / " , " ! " ] ) )
async def start ( client , message ) :
2023-03-18 00:04:54 +02:00
# Photo = "https://i.imgur.com/79hHVX6.png"
2023-01-12 06:37:56 +02:00
2023-04-22 22:18:10 +03:00
buttons = [
[
InlineKeyboardButton (
" Add to your group " , url = " https://t.me/gootmornbot?startgroup=true "
)
]
]
2023-03-18 00:04:54 +02:00
await message . reply_text (
# photo=Photo,
2023-04-22 22:18:10 +03:00
text = f " Hello! \n ask me to imagine anything \n \n /draw text to image " ,
reply_markup = InlineKeyboardMarkup ( buttons ) ,
)
2023-01-12 06:37:56 +02:00
app . run ( )