ro/stable-diffusion-telegram-bot
1
0
Fork 0
forked from yair/stable-diffusion-telegram-bot
Code Pull Requests Activity

ADD choose model

Browse Source
This commit is contained in:
yair 2023-03-18 00:04:54 +02:00
parent 1306780426
commit cae44e0e3b
2 changed files with 87 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@ -1,2 +1,3 @@
*.png
.env
*.session

116
main.py
View File

@ -1,30 +1,41 @@
from pyrogram import Client, filters
from pyrogram.types import *
import os
import json
import requests
import io
import os
import random
from PIL import Image, PngImagePlugin
import base64
from PIL import Image, PngImagePlugin
from pyrogram import Client, filters
from pyrogram.types import *
# Done! Congratulations on your new bot. You will find it at
# 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
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("TOKEN", None)
SD_URL = os.environ.get("SD_URL", None)
print(SD_URL)
bot = Client(
app = Client(
"stable",
api_id=API_ID,
api_hash=API_HASH,
bot_token=TOKEN
)
@app.on_message(filters.command(["draw"]))
def draw(client, message):
msgs = message.text.split(' ', 1)
if len(msgs) == 1:
message.reply_text("Format : /draw < text to anime image >")
message.reply_text("Format : /draw < text to image >")
return
msg = msgs[1]
@ -34,25 +45,29 @@ def draw(client, message):
r = requests.post(url=f'{SD_URL}/sdapi/v1/txt2img', json=payload).json()
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)
word = f"{message.from_user.id}-MOE{gen1}{gen2}{gen3}{gen4}{gen5}{gen6}{gen7}{gen8}{gen9}{gen10}"
# save the image to a file
def genr(): # generate random string
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}"
word=genr()
for i in r['images']:
image = Image.open(io.BytesIO(base64.b64decode(i.split(",", 1)[0])))
png_payload = {"image": "data:image/png;base64," + i}
response2 = requests.post(url=f'{url}/sdapi/v1/png-info',
response2 = requests.post(url=f'{SD_URL}/sdapi/v1/png-info',
json=png_payload)
pnginfo = PngImagePlugin.PngInfo()
@ -62,27 +77,66 @@ def draw(client, message):
message.reply_photo(
photo=f"{word}.png",
caption=
f"Prompt - **{msg}**\n **[{message.from_user.first_name}-Kun](tg://user?id={message.from_user.id})**\n Join @WaifuAiSupport"
f"Prompt - **{msg}**\n **[{message.from_user.first_name}-Kun](tg://user?id={message.from_user.id})**"
)
os.remove(f"{word}.png")
K.delete()
# get the models, choose one, and set it as the model checkpoint
@app.on_message(filters.command(["getmodels"]))
async def get_models(client, message):
response = requests.get(url=f'{SD_URL}/sdapi/v1/sd-models')
if response.status_code == 200:
models_json = response.json()
# create buttons for each model name
buttons = []
for model in models_json:
buttons.append([InlineKeyboardButton(model['title'], callback_data=model['model_name'])])
# send the message
await message.reply_text(
text="Select a model [checkpoint] to use",
reply_markup=InlineKeyboardMarkup(buttons)
)
@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
options = {
"sd_model_checkpoint": sd_model_checkpoint
}
# post the options
response = requests.post(url=f'{SD_URL}/sdapi/v1/options', json=options)
if response.status_code == 200:
# if the post was successful, send a message
await callback_query.message.reply_text("checpoint set to " + sd_model_checkpoint)
else:
# if the post was unsuccessful, send an error message
await callback_query.message.reply_text("Error setting options")
# start the bot
@app.on_message(filters.command(["start"], prefixes=["/", "!"]))
async def start(client, message):
Photo = "https://media.discordapp.net/attachments/1028156834944655380/1062018608022171788/3aac7aaf-0065-40aa-9e4d-430c717b3d87.jpg"
# Photo = "https://i.imgur.com/79hHVX6.png"
buttons = [[
InlineKeyboardButton("Add to your group",
url="http://t.me/botname?startgroup=true"),
InlineKeyboardButton("Channel", url="https://t.me/otakatsu"),
InlineKeyboardButton("Support", url="https://t.me/otakatsu_chat")
url="https://t.me/gootmornbot?startgroup=true")
]]
await message.reply_photo(
photo=Photo,
caption=
f"Hello! I'm botname Ai and I can make an anime-styled picture!\n\n/generate - Reply to Image\n/draw text to anime image\n\nPowered by @Otakatsu",
reply_markup=InlineKeyboardMarkup(buttons))
await message.reply_text(
# photo=Photo,
text=
f"Hello!\nask me to imagine anything\n\n/draw text to image",
reply_markup=InlineKeyboardMarkup(buttons)
)
app.run()