Initial commit

This commit is contained in:
Ed 2022-11-17 23:44:39 +01:00
commit 47ce04e15f
12 changed files with 480 additions and 0 deletions

28
bot.py Normal file
View File

@ -0,0 +1,28 @@
import logging, asyncio
from telethon import events
from telethon.tl.custom import Button
from os.path import isfile
from time import time
import httpx
import re
import sqlite3
from os import unlink
from random import randint
from telethon.errors.rpcerrorlist import MessageNotModifiedError, UserIsBlockedError
from telethon.utils import get_display_name
from config import *
from judge_prompt import judge
import handlers.welcome
import handlers.info
import handlers.prompt
import handlers.help
if __name__ == '__main__':
for x in handlers.welcome.handler: client.add_event_handler(x)
for x in handlers.info.handler: client.add_event_handler(x)
for x in handlers.prompt.handler: client.add_event_handler(x)
for x in handlers.help.handler: client.add_event_handler(x)
client.start()
client.flood_sleep_threshold = 24*60*60
client.run_until_disconnected()

Binary file not shown.

Binary file not shown.

Binary file not shown.

Binary file not shown.

12
handlers/analyze.py Normal file
View File

@ -0,0 +1,12 @@
from telethon import events
@events.register(events.callbackquery.CallbackQuery(pattern=r'^analyze$'))
async def analyze_prompt(ev):
log.info(f'{(ev.input_sender.user_id, get_display_name(ev.sender))}: analyze')
res = conn.execute('SELECT * FROM pending_prompt WHERE user_id = ? LIMIT 1', (ev.input_sender.user_id,)).fetchone()
if res:
comments, quality = judge(res['prompt'])
await ev.respond("\n".join(comments), parse_mode='HTML')
else:
await ev.respond('What am i supposed to analyze?')

27
handlers/help.py Normal file
View File

@ -0,0 +1,27 @@
from telethon import events
from telethon.utils import get_display_name
from telethon.events import StopPropagation
from telethon.tl.custom import Button
import logging
log = logging.getLogger('welcome')
@events.register(events.NewMessage(pattern='^/help', incoming=True))
async def help(ev):
await ev.respond(f"""Full list of commands:\n
/help - get this message
/new - create a new prompt
/edit - edit your current prompt
/copy - copy your last prompt
/start - get the welcome message
/info - get info about the current status of the bot
/register - register your Horde key
/kudos - details about your kudos
/models - get the list of models
""")
@events.register(events.NewMessage(pattern='^/register$', incoming=True))
async def register(ev):
await ev.respond("You can register your own Horde api key by writing `/register YOUR_HORDE_KEY_HERE`.\nBy registering your own account, you will stop being a guest and you will be able to send and receive kudos from/to users, unlock higher limits and earn kudos by sharing your GPU.", buttons=[Button.url('Register to the Horde', 'https://stablehorde.net/register')])
handler = [help,register]

52
handlers/info.py Normal file
View File

@ -0,0 +1,52 @@
from telethon import events
from telethon.utils import get_display_name
from telethon.events import StopPropagation
import logging
import httpx
from config import *
log = logging.getLogger('info')
@events.register(events.NewMessage(pattern='^/info', incoming=True, func=lambda e: e.is_private))
async def info(ev):
try:
async with httpx.AsyncClient() as client:
res = await client.get("https://stablehorde.net/api/v2/find_user", headers={'apikey': api_horde})
data = res.json()
except:
await ev.respond('There was an issue retrieving details. Try later')
else:
await ev.respond(f"""<strong>{data['username']}'s <a href="https://stablehorde.net/">stable horde</a></strong>\nKudos: {data['kudos']}\nOnline workers: {data['worker_count']}\nTotal images: {data['contributions']['fulfillments']}\nTotal megapixelsteps: {data['contributions']['megapixelsteps']}""", parse_mode='html')
@events.register(events.NewMessage(pattern='^/kudos', incoming=True, func=lambda e: e.is_private))
async def kudos(ev):
apikey = conn.execute('SELECT key FROM user WHERE id = ? AND key IS NOT NULL', (ev.input_sender.user_id,)).fetchone()
try:
async with httpx.AsyncClient() as client:
res = await client.get("https://stablehorde.net/api/v2/find_user", headers={'apikey': default_horde_key if not apikey else apikey[0]})
data = res.json()
except:
await ev.respond('There was an issue retrieving details. Try later')
raise
else:
await ev.respond(f"""Hello {data['username']}, you have {data['kudos']} kudos.\nA higher amount of kudos will lead to faster generations!""" + ('\n\n⚠️ You are using the bot as a guest: register on <a href="https://stablehorde.net">Stable Horde</a> and donate GPU time to have higher priority' if not apikey else ''), parse_mode='html')
@events.register(events.NewMessage(pattern='^/models', incoming=True, func=lambda e: e.is_private))
async def models(ev):
try:
async with httpx.AsyncClient() as client:
res = await client.get("https://stablehorde.net/api/v2/status/models")
data = res.json()
except:
await ev.respond('There was an issue retrieving details. Try later')
raise
else:
ret = "Available models (number of workers):"
for mod in data:
if mod['name'] not in enabled_models: continue
ret += f"\n- {mod['name']} ({mod['count']})"
await ev.respond(ret, parse_mode='html')
handler = [info,kudos,models]

214
handlers/prompt.py Normal file
View File

@ -0,0 +1,214 @@
from telethon import events
from telethon.utils import get_display_name
from telethon.events import StopPropagation, CallbackQuery
from telethon.tl.custom import Button
from telethon.errors import MessageNotModifiedError
import logging
import httpx
import json
from random import randint
from sqlite3 import IntegrityError
from config import *
from parameters import *
from hashlib import md5
log = logging.getLogger('prompt')
def get_prompt(user_id, delete=True):
prompt = conn.execute('SELECT payload FROM pending_prompt WHERE user_id = ?', (user_id,)).fetchone()
if prompt:
prompt = json.loads(prompt[0])
else:
return None
if prompt['prompt']: return prompt
if delete:
conn.execute('DELETE FROM pending_prompt WHERE user_id = ?', (user_id,))
return None
else:
return prompt
@events.register(events.callbackquery.CallbackQuery(pattern=r'^new_prompt'))
@events.register(events.NewMessage(pattern='^/new', incoming=True, func=lambda e: e.is_private))
async def new_prompt(ev):
prompt = get_prompt(ev.input_sender.user_id)
if prompt:
await ev.respond("You already have a pending prompt!",
buttons=[[Button.inline(f"Delete", f"delete_prompt"), Button.inline(f"Edit", "edit_prompt")]])
raise StopPropagation
req_msg = await ev.respond("A new prompt! What would you like to generate?\nSend a message with a phrase or a list of tags you would like to generate. For ideas, check @ai621gen or ask in our chat @ai621chat.")
conn.execute('INSERT INTO pending_prompt(user_id, payload) VALUES (?,?)', (ev.input_sender.user_id, json.dumps(fields_template)))
conn.execute('UPDATE user SET pending = \'prompt\', pending_msg = ? WHERE id = ?', (req_msg.id, ev.input_sender.user_id))
print('Has been created')
raise StopPropagation
@events.register(events.NewMessage(pattern='^/delete', incoming=True, func=lambda e: e.is_private))
@events.register(events.callbackquery.CallbackQuery(pattern=r'^delete_prompt'))
async def delete_prompt(ev):
conn.execute('DELETE FROM pending_prompt WHERE user_id = ?', (ev.input_sender.user_id,))
await ev.respond('Your pending prompt, if any, has been deleted.',
buttons=[[Button.inline(f"New prompt", f"new_prompt")]])
raise StopPropagation
@events.register(events.callbackquery.CallbackQuery(pattern=r'^msg_but:'))
@events.register(events.NewMessage(incoming=True, pattern='^([^\/](\n|.)*)?$', func=lambda e: e.is_private))
async def accept_data(ev):
pending = conn.execute('SELECT pending FROM user WHERE id = ? AND pending IS NOT NULL', (ev.input_sender.user_id,)).fetchone()
if not pending:
await ev.respond("I don't understand. Maybe try to /start again?")
raise StopPropagation
else:
pending = pending[0]
print(pending)
prompt = get_prompt(ev.input_sender.user_id, delete=False)
if not prompt:
await ev.respond('You have no pending prompt to edit. Try to /start again?')
raise StopPropagation
if hasattr(ev, 'message'):
data = ev.message.raw_text.strip()
else:
data = ev.data.decode().split(':', 1)[1]
await ev.respond(f'✅ You have set {pending} to {data}')
if pending == 'prompt':
prompt['prompt'] = data
conn.execute('UPDATE pending_prompt SET payload = ? WHERE user_id = ?', (json.dumps(prompt), ev.input_sender.user_id))
conn.execute('UPDATE user SET pending = NULL WHERE id = ?', (ev.input_sender.user_id,))
await edit_prompt(ev)
raise StopPropagation
@events.register(events.callbackquery.CallbackQuery(pattern=r'^(change|toggle):[a-z]+$'))
async def edit_parameter(ev):
prompt = get_prompt(ev.input_sender.user_id)
if not prompt:
await ev.delete()
await ev.respond('You have no pending prompt to edit.', buttons=[[Button.inline(f"New prompt", f"new_prompt")]])
raise StopPropagation
action, parameter = ev.data.decode().split(':', 1)
if action == 'toggle':
if limits[parameter]['type'] == 'boolean':
prompt[parameter] = not prompt[parameter]
conn.execute('UPDATE pending_prompt SET payload = ? WHERE user_id = ?', (json.dumps(prompt), ev.input_sender.user_id))
await edit_prompt(ev)
if action == 'change':
btns = [Button.inline(str(x), 'msg_but:'+str(x)) for x in field_buttons.get(parameter, [])] or None
print(btns)
await ev.respond(prompt_msg[parameter], buttons=btns)
conn.execute('UPDATE user SET pending = ? WHERE id = ?', (parameter, ev.input_sender.user_id))
@events.register(events.callbackquery.CallbackQuery(pattern=r'^edit_prompt'))
@events.register(events.NewMessage(pattern='^/edit', incoming=True, func=lambda e: e.is_private))
async def edit_prompt(ev):
prompt = get_prompt(ev.input_sender.user_id)
if not prompt:
await ev.delete()
await ev.respond('You have no pending prompt to edit.', buttons=[[Button.inline(f"New prompt", f"new_prompt")]])
raise StopPropagation
msg_id = conn.execute('SELECT msg_id FROM pending_prompt WHERE user_id = ?', (ev.input_sender.user_id,)).fetchone()[0]
message = (f"<strong>Your current prompt.</strong>\n\n"
f"📣 Prompt: {prompt['prompt']}\n"
f"🚫 Negative prompt: {prompt['negative_prompt']}\n\n"
f"<strong>Generation parameters:</strong>\n"
f"🎛 Model: {prompt['models']}\n"
f"🌱 Seed: {prompt['seed']}\n"
f"♨️ Sampler: {prompt['sampler_name']}\n"
f"💎 Config scale: {prompt['cfg_scale']}\n"
f"🌀 Steps: {prompt['steps']}\n\n"
f"<strong>Output options:</strong>\n"
f"🖥 Resolution: {prompt['resolution']}\n"
f"✨ Upscaling: {prompt['use_upscaling']}\n"
f"🔢 Number of images: {prompt['n']}\n\n"
+ (
(
f"<strong>img2img options:</strong>\n"
f"📷 Base image: {'uploaded image' if prompt['image'] else 'no image'}\n",
) if prompt['image'] else ''
))
kwargs = {
'parse_mode': 'html',
'link_preview': False,
'buttons': [
[
Button.inline(f"👀 Prompt", f"change:prompt"),
Button.inline(f"⛔ Negative prompt", f"change:negative_prompt"),
],
[
Button.inline(f"🌱 Seed", f"change:seed"),
Button.inline(f"🔢 Number", f"change:n"),
Button.inline(f"✨ Upscale", f"toggle:use_upscaling")
],
[
Button.inline(f"💎 Config scale", f"change:cfg_scale"),
Button.inline(f"🌀 Steps", f"change:steps"),
],
[
# Button.inline(f"📷 Base image", f"change:image"),
Button.inline(f"🖥 Resolution", f"change:resolution"),
],
[
Button.inline(f"✅ Confirm", f"confirm_prompt"),
Button.inline(f"🕵️ Analyze", f"analyze_prompt"),
Button.inline(f"❌ Delete", f"delete_prompt")
]
]
}
if msg_id:
try:
await client.edit_message(ev.input_sender, msg_id, message, **kwargs)
except MessageNotModifiedError:
return
except Exception as e:
log.error(str(e))
else:
return
log.error(f"There was an issue editing the message of pending prompt")
msg_id = await ev.respond(message, **kwargs)
conn.execute('UPDATE pending_prompt SET msg_id = ? WHERE user_id = ?', (msg_id.id, ev.input_sender.user_id))
@events.register(events.callbackquery.CallbackQuery(pattern=r'^confirm_prompt'))
@events.register(events.NewMessage(pattern='^/confirm', incoming=True, func=lambda e: e.is_private))
async def confirm_prompt(ev):
prompt = get_prompt(ev.input_sender.user_id)
if not prompt:
await ev.respond('You have no prompt to confirm!', buttons=[[Button.inline(f"New prompt", f"new_prompt")]])
raise StopPropagation
json_prompt = prompt_template.copy()
json_prompt['prompt'] = prompt['prompt']
json_prompt['params']['width'], json_prompt['params']['height'] = prompt['resolution'].split('x')
if prompt['seed'] == 'random': prompt['seed'] = randint(0, 1000000)
if prompt['negative_prompt']:
json_prompt['prompt'].append('### ' + prompt['negative_prompt'])
for n in ['sampler_name', 'cfg_scale', 'seed', 'use_upscaling', 'steps', 'n', 'models']:
json_prompt['params'][n] = prompt[n]
conn.execute('INSERT INTO prompt(user_id, original_payload, payload, hash) VALUES (?,?,?,?)',
(ev.input_sender.user_id, json.dumps(prompt), json.dumps(json_prompt), md5(json.dumps(json_prompt, sort_keys=True)).hexdigest())
)
handler = [new_prompt, accept_data, delete_prompt, confirm_prompt, edit_prompt, edit_parameter]

29
handlers/welcome.py Normal file
View File

@ -0,0 +1,29 @@
from telethon import events
from telethon.utils import get_display_name
from telethon.events import StopPropagation
from telethon.tl.custom import Button
import logging
from config import *
log = logging.getLogger('welcome')
@events.register(events.NewMessage(pattern='^/start', incoming=True, func=lambda e: e.is_private))
async def welcome(ev):
exists = conn.execute('SELECT 1 FROM user WHERE id = ?', (ev.input_sender.user_id,))
if not exists.fetchone():
conn.execute('INSERT INTO user(id, first_name, last_name, lang_code) VALUES (?,?,?,?)',
(ev.input_sender.user_id, ev.sender.first_name, ev.sender.last_name, ev.sender.lang_code))
log.info(f'{(ev.input_sender.user_id, get_display_name(ev.sender))}: hello')
await ev.respond(f'Hello, and welcome to ai621. The guidelines are simple:\n1. No prompts that could result in abusive images or underage characters\nThis bot is powered by <strong>Stable Horde</strong>.', buttons=[[Button.url('👥 Group (NSFW)', 'https://t.me/ai621chat'),], [Button.url('📣 News & Best images (NSFW)', 'https://t.me/ai621gen'),],[Button.url('🔗 Join the Horde', 'https://stablehorde.net'),],[Button.inline('🆕 New prompt', 'new_prompt')]], parse_mode='html')
@events.register(events.callbackquery.CallbackQuery())
@events.register(events.NewMessage(incoming=True, func=lambda e: e.is_private))
async def precheck(ev):
exists = conn.execute('SELECT 1 FROM user WHERE id = ?', (ev.input_sender.user_id,))
if not exists.fetchone():
log.info(f'{(ev.input_sender.user_id, get_display_name(ev.sender))}: invite start')
await ev.client.send_message(ev.input_sender, 'Your profile is incomplete. Please use /start again.')
raise StopPropagation
handler = [welcome, precheck]

34
parameters.py Normal file
View File

@ -0,0 +1,34 @@
enabled_fields = ['prompt', 'sampler_name', 'cfg_scale', 'resolution', 'use_upscaling', 'seed', 'steps', 'n', 'models']
field_buttons = {
'sampler_name': ['k_lms', 'k_heun', 'k_euler', 'k_euler_a', 'k_dpm_2', 'k_dpm_2_a', 'k_dpm_fast', 'k_dpm_adaptive', 'k_dpmpp_2s_a', 'k_dpmpp_2m'],
'cfg_scale': ['3.0', '5.0', '9.0', '15.0'],
'resolution': ['768x512', '1024x512', '512x512', '512x1024', '512x768'],
'use_upscaling': ['True', 'False'],
'seed': ['random',],
'steps': ['20', '30', '50', '75', '100'],
'n': ['1', '2', '3', '4'],
'models': ['Yiffy', 'Furry Epoch', 'Zack3D']
}
limits = {
'cfg_scale': {'type': 'float', 'min': 1, 'max': 20, 'multiple': 0.5},
'use_upscaling': {'type': 'boolean'},
'steps': {'type': 'integer', 'min': 5, 'max': 100, 'multiple': 5},
'n': {'type': 'integer', 'min': 1, 'max': 5},
'models': {'type': 'enum', 'allowed': ['Yiffy', 'Furry Epoch', 'Zack3D']},
'seed': {'type': 'integer', 'min': 0, 'max': 1000000}
}
prompt_msg = {
'prompt': 'Tell me the prompt. Remember to consider the model you want to use.',
'negative_prompt': 'Tell me the negative prompt. This is what you <strong>don\'t</strong> want to see in the final image.',
'sampler_name': 'What sampler do you want? Sampler changes the way noise is transformed into images.',
'cfg_scale': 'The config scale specifies the "hardness" of the neural network: higher values will create stronger shapes. Please write a number from 1 to 50 or press a button.',
'resolution': 'The resolution specifies the size and format of the image generated. Choose one, or write your own (up to 3072x3072, multiples of 64).',
'seed': 'The seed specifies the unique shape of the generated image(s). Write a number from 0 to 1000000, or click on "random".',
'steps': 'The amount of steps specified how much time will be spent generating the image. Higher values usually generate more defined images, but your wait time will dramatically increase! Choose one or write a number between 10 and 100.',
'n': 'How many images do you want to generate? Up to 4.',
'models': 'Please choose a model to use!\n- <strong>Yiffy</strong> = trained on top e621 posts\n- <strong>Furry Epoch</strong> = trained on 300k images from e621\n- <strong>Zack3D</strong> = trained on 100k images from e621, specializes in transformation, latex, tentacles, goo, ferals and bondage\n\nYou can also give me more than one prompt, split by commas, to generate multiple images with different models.'
}

84
res.py Normal file
View File

@ -0,0 +1,84 @@
from random import randint
from telethon.tl.custom import Button
params = {
'seed': {
'description': 'The seed is a number between 0 and 1000000. The seed defines the "shape" of the noise which will used to create the image(s).',
'default': (lambda: randint(0, 1000000)),
'transform': lambda x: int(x),
'checks': {
(lambda x: x < 1000000): 'The number must be less than 1000000',
(lambda x: x > 0): 'The number must be more than 0',
}
},
'image': {
'description': 'Upload an image, give me an e621 link or just delete the current one.',
'default': None,
'buttons': [Button.inline(f"🚫 No image", f"msg_but:no_image")]
},
'prompt': {
'description': 'Give me a list of e621 tags, a phrase or an e621 link to use as a prompt for your image. You can check which tags you can use over at https://ai621.foxo.me/test.html',
'default': None,
'transform': lambda x: x.strip(),
'checks': {
(lambda x: len(x) < 200): 'The prompt cannot be longer than 200 characters.',
(lambda x: len(x) > 32): 'The prompt cannot be shorter than 32 characters.',
(lambda x: ('🖼' in x) or ('👀' in x) or ('🌀' in x)): 'If you copy and paste other prompts, try to clean them from extra emojis and formatting.',
}
},
'negative_prompt': {
'description': 'This is the negative prompt. Use this to define what you <strong>don\'t</strong> want to see. ',
'default': None,
'transform': lambda x: x.strip(),
'checks': {
(lambda x: len(x) < 200): 'The neg prompt cannot be longer than 200 characters.',
(lambda x: len(x) > 16): 'The neg prompt cannot be shorter than 16 characters.',
(lambda x: ('🖼' in x) or ('👀' in x) or ('🌀' in x)): 'If you copy and paste other prompts, try to clean them from extra emojis and formatting.',
}
},
'number': {
'description': 'How many images would you like to receive?',
'default': 3,
'transform': lambda x: int(3),
'checks': {
(lambda x: x in [1,2,3,4]): 'You can only generate between 1 and 4 images.',
},
'buttons': [Button.inline(f"1", f"msg_but:1"), Button.inline(f"2", f"msg_but:2"), Button.inline(f"3", f"msg_but:3"), Button.inline(f"4", f"msg_but:4")],
},
'detail': {
'description': 'Detail (guidance scale) how heavily should the bot draw tags? This value should be between 2 and 20. (and a multiple of 0.1)',
'default': 6.9,
'transform': lambda x: round(float(x), 1),
'checks': {
(lambda x: len(x) < 20): 'The maximum for this value is 20.',
(lambda x: len(x) > 2): 'The minimum for this value is 2.',
},
'buttons': [Button.inline(f"S (3.0)", f"msg_but:3"), Button.inline(f"M (6.5)", f"msg_but:6.5"), Button.inline(f"L (9.0)", f"msg_but:9")],
},
'inference_steps': {
'description': 'Cycles (inference steps) how many times should we try to redraw the image before giving the result? Higher values = higher quality, but only when prompts are good.',
'default': 40,
'transform': lambda x: int(x),
'checks': {
(lambda x: len(x) < 200): 'The maximum for this value is 200.',
(lambda x: len(x) > 20): 'The minimum for this value is 20.',
},
'buttons': [Button.inline(f"S (20)", f"msg_but:20"), Button.inline(f"M (40)", f"msg_but:40"), Button.inline(f"L (60)", f"msg_but:60")],
},
'resolution': {
'description': 'Decide the aspect ratio and resolution of the final image.',
'default': '512x512',
'checks': {
(lambda x: x in ['512x512', '512x768', '768x512', '1024x512', '512x1024']): 'This resolution is invalid.'
},
'buttons': [[Button.inline(f"Potrait (512x768)", f"msg_but:512x768"), Button.inline(f"Landscape (768x512)", f"msg_but:768x512")],[Button.inline(f"Square (512x512)", f"msg_but:512x512")],[Button.inline(f"Ultrawide (1024x512)", f"msg_but:1024x512"), Button.inline(f"Ultratall (512x1024)", f"msg_but:512x1024")]]
},
'resolution': 'The width and height of the final image.',
'crop': 'Do you want to crop the base image so that it matches the generated image resolution?',
'hires': 'Do you want to receive a high res image at the end of the generation? (it will take more time)',
}
# 'blend': [Button.inline(f"S (0.3)", f"msg_but:0.3"), Button.inline(f"M (0.6)", f"msg_but:0.6"), Button.inline(f"L (0.8)", f"msg_but:0.8")],