yair/stable-diffusion-telegram-bot
2
1
Fork 1
mirror of https://github.com/5shekel/stable-diffusion-telegram-bot.git synced 2024-05-22 19:33:14 +03:00
Code Issues Packages Projects Releases Wiki Activity

format

Browse Source
This commit is contained in:
yair 2023-04-22 22:18:10 +03:00
commit b038ef5573
1 changed files with 60 additions and 64 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

110
main.py
View File

@ -24,48 +24,40 @@ API_HASH = os.environ.get("API_HASH", None)
TOKEN = os.environ.get("TOKEN", None) TOKEN = os.environ.get("TOKEN", None)
SD_URL = os.environ.get("SD_URL", None) SD_URL = os.environ.get("SD_URL", None)
print(SD_URL) print(SD_URL)
app = Client( app = Client("stable", api_id=API_ID, api_hash=API_HASH, bot_token=TOKEN)
"stable",
api_id=API_ID,
api_hash=API_HASH, def slice_positive_negative(string):
bot_token=TOKEN delimiter = "ng"
) if delimiter in string:
index = string.index(delimiter)
positive = string[:index].rstrip()
negative = string[index + len(delimiter) :].lstrip()
return positive, negative
else:
return string, "null" # if we are missing the ng bad thing happen....
@app.on_message(filters.command(["draw"])) @app.on_message(filters.command(["draw"]))
def draw(client, message): def draw(client, message):
msgs = message.text.split(' ', 1) msgs = message.text.split(" ", 1)
if len(msgs) == 1: if len(msgs) == 1:
message.reply_text("Format : /draw < text to image >") message.reply_text(
"Format :\n/draw < text to image >\nng < negative (optional) >"
)
return return
msg = msgs[1]
K = message.reply_text("Please Wait 10-15 Second") msg = slice_positive_negative(msgs[1])
print(msg)
prompt = msg payload = {
negative_prompt = "lowres, low resolution, bad resolution, bad, poor quality, bad quality, blurry, bad art, incomplete, logo, signature, text, jpeg artifacts, compression artifacts,child, childish" "prompt": msg[0],
"negative_prompt": msg[1],
# Look for negative prompt in brackets after "NG" in the message }
ng_index = msg.find("ng" or "NG")
if ng_index != -1:
bracket_start_index = msg.find("[", ng_index)
bracket_end_index = msg.find("]", ng_index)
if bracket_start_index != -1 and bracket_end_index != -1:
new_negative_prompt = msg[bracket_start_index+1:bracket_end_index]
if negative_prompt == "":
negative_prompt = new_negative_prompt
else:
negative_prompt = f"{negative_prompt}, {new_negative_prompt}"
prompt = msg[:ng_index] + msg[bracket_end_index+1:]
else:
prompt = msg[:ng_index]
payload = {"prompt": prompt}
if negative_prompt != "":
payload["negative_prompt"] = negative_prompt
print(payload) print(payload)
r = requests.post(url=f'{SD_URL}/sdapi/v1/txt2img', json=payload).json()
K = message.reply_text("Please Wait 10-15 Second")
r = requests.post(url=f"{SD_URL}/sdapi/v1/txt2img", json=payload).json()
def genr(): def genr():
chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
@ -83,42 +75,46 @@ def draw(client, message):
return f"{message.from_user.id}-MOE{gen1}{gen2}{gen3}{gen4}{gen5}{gen6}{gen7}{gen8}{gen9}{gen10}" return f"{message.from_user.id}-MOE{gen1}{gen2}{gen3}{gen4}{gen5}{gen6}{gen7}{gen8}{gen9}{gen10}"
word = genr() word = genr()
for i in r['images']: for i in r["images"]:
image = Image.open(io.BytesIO(base64.b64decode(i.split(",", 1)[0]))) image = Image.open(io.BytesIO(base64.b64decode(i.split(",", 1)[0])))
png_payload = {"image": "data:image/png;base64," + i} png_payload = {"image": "data:image/png;base64," + i}
response2 = requests.post(url=f'{SD_URL}/sdapi/v1/png-info', response2 = requests.post(url=f"{SD_URL}/sdapi/v1/png-info", json=png_payload)
json=png_payload)
pnginfo = PngImagePlugin.PngInfo() pnginfo = PngImagePlugin.PngInfo()
pnginfo.add_text("parameters", response2.json().get("info")) pnginfo.add_text("parameters", response2.json().get("info"))
image.save(f'{word}.png', pnginfo=pnginfo) image.save(f"{word}.png", pnginfo=pnginfo)
message.reply_photo( message.reply_photo(
photo=f"{word}.png", photo=f"{word}.png",
caption= caption=f"Prompt - **{msg[0]}**\nNegative Prompt - **{msg[1]}**\n**[{message.from_user.first_name}-Kun](tg://user?id={message.from_user.id})**",
f"Prompt - **{msg}**\n **[{message.from_user.first_name}-Kun](tg://user?id={message.from_user.id})**"
) )
# os.remove(f"{word}.png") # os.remove(f"{word}.png")
K.delete() K.delete()
@app.on_message(filters.command(["getmodels"])) @app.on_message(filters.command(["getmodels"]))
async def get_models(client, message): async def get_models(client, message):
response = requests.get(url=f'{SD_URL}/sdapi/v1/sd-models') response = requests.get(url=f"{SD_URL}/sdapi/v1/sd-models")
if response.status_code == 200: if response.status_code == 200:
models_json = response.json() models_json = response.json()
# create buttons for each model name # create buttons for each model name
buttons = [] buttons = []
for model in models_json: for model in models_json:
buttons.append([InlineKeyboardButton(model['title'], callback_data=model['model_name'])]) buttons.append(
[
InlineKeyboardButton(
model["title"], callback_data=model["model_name"]
)
]
)
# send the message # send the message
await message.reply_text( await message.reply_text(
text="Select a model [checkpoint] to use", text="Select a model [checkpoint] to use",
reply_markup=InlineKeyboardMarkup(buttons) reply_markup=InlineKeyboardMarkup(buttons),
) )
@app.on_callback_query() @app.on_callback_query()
async def process_callback(client, 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 # if a model button is clicked, set sd_model_checkpoint to the selected model's title
@ -127,35 +123,35 @@ async def process_callback(client, callback_query):
# The sd_model_checkpoint needs to be set to the title from /sdapi/v1/sd-models # The sd_model_checkpoint needs to be set to the title from /sdapi/v1/sd-models
# post using /sdapi/v1/options # post using /sdapi/v1/options
options = { options = {"sd_model_checkpoint": sd_model_checkpoint}
"sd_model_checkpoint": sd_model_checkpoint
}
# post the options # post the options
response = requests.post(url=f'{SD_URL}/sdapi/v1/options', json=options) response = requests.post(url=f"{SD_URL}/sdapi/v1/options", json=options)
if response.status_code == 200: if response.status_code == 200:
# if the post was successful, send a message # if the post was successful, send a message
await callback_query.message.reply_text("checpoint set to " + sd_model_checkpoint) await callback_query.message.reply_text(
"checpoint set to " + sd_model_checkpoint
)
else: else:
# if the post was unsuccessful, send an error message # if the post was unsuccessful, send an error message
await callback_query.message.reply_text("Error setting options") await callback_query.message.reply_text("Error setting options")
@app.on_message(filters.command(["start"], prefixes=["/", "!"])) @app.on_message(filters.command(["start"], prefixes=["/", "!"]))
async def start(client, message): async def start(client, message):
# Photo = "https://i.imgur.com/79hHVX6.png" # Photo = "https://i.imgur.com/79hHVX6.png"
buttons = [[ buttons = [
InlineKeyboardButton("Add to your group", [
url="https://t.me/gootmornbot?startgroup=true") InlineKeyboardButton(
]] "Add to your group", url="https://t.me/gootmornbot?startgroup=true"
)
]
]
await message.reply_text( await message.reply_text(
# photo=Photo, # photo=Photo,
text= text=f"Hello!\nask me to imagine anything\n\n/draw text to image",
f"Hello!\nask me to imagine anything\n\n/draw text to image", reply_markup=InlineKeyboardMarkup(buttons),
reply_markup=InlineKeyboardMarkup(buttons)
) )