mirror of
https://git.wownero.com/dsc/craiyon-irc.git
synced 2024-08-15 01:03:24 +00:00
244 lines
7.4 KiB
Python
244 lines
7.4 KiB
Python
import re
|
|
import time
|
|
import os
|
|
import random
|
|
import json
|
|
from base64 import b64decode
|
|
from typing import Optional, List
|
|
from datetime import datetime
|
|
import logging
|
|
import asyncio
|
|
from io import BytesIO
|
|
import bottom
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
from PIL.JpegImagePlugin import JpegImageFile
|
|
import aiohttp
|
|
|
|
from settings import *
|
|
|
|
TASK_QUEUE = asyncio.Queue()
|
|
TASK_AUTHORS = {}
|
|
FONT = ImageFont.truetype('font-bold.ttf', 22)
|
|
bot = bottom.Client(host=IRC_HOST, port=IRC_PORT, ssl=IRC_TLS, loop=asyncio.get_event_loop())
|
|
|
|
USER_AGENT = 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/97.0.4692.99 Safari/537.36'
|
|
HEADERS = {
|
|
'User-Agent': USER_AGENT,
|
|
'Accept': 'application/json',
|
|
'Accept-Language': 'en-US,en;q=0.5',
|
|
'Referer': 'https://www.craiyon.com/',
|
|
'Origin': 'https://www.craiyon.com',
|
|
'Connection': 'keep-alive',
|
|
'Sec-Fetch-Dest': 'empty',
|
|
'Sec-Fetch-Mode': 'cors',
|
|
'Sec-Fetch-Site': 'same-site',
|
|
'Pragma': 'no-cache',
|
|
'Cache-Control': 'no-cache',
|
|
}
|
|
|
|
|
|
class Task:
|
|
def __init__(self, author: str, channel: str, term: str, dt: int):
|
|
self.author = author
|
|
self.channel = channel
|
|
self.term = term
|
|
self.dt = dt
|
|
|
|
async def process_images(self, images: List[JpegImageFile]) -> JpegImageFile:
|
|
if len(images) <= 0:
|
|
raise Exception("Not enough images")
|
|
|
|
if MODE == "monero":
|
|
bg = Image.open('bg-pools.jpg')
|
|
term_max = 60
|
|
term_x_start = 100
|
|
if MODE == "2600nl":
|
|
bg = Image.open('bg-2600nl.jpg')
|
|
term_max = 52
|
|
term_x_start = 230
|
|
else:
|
|
bg = Image.open('bg-wownero.jpg')
|
|
term_max = 52
|
|
term_x_start = 230
|
|
|
|
d1 = ImageDraw.Draw(bg)
|
|
|
|
term = self.term
|
|
if len(term) >= term_max:
|
|
term = self.term[:term_max] + ".."
|
|
|
|
d1.text((term_x_start, 160), term, font=FONT, fill=(255, 255, 255))
|
|
|
|
x_start = 41
|
|
y_start = 232
|
|
for i, img in enumerate(images):
|
|
x_offset = (i % 3) * (img.width + 20)
|
|
y_offset = int(i / 3) * (img.width + 20)
|
|
bg.paste(img, (x_start + x_offset, y_start + y_offset))
|
|
return bg
|
|
|
|
async def post_image(self, image: BytesIO) -> str:
|
|
url = 'https://uguu.se/upload.php'
|
|
timeout = aiohttp.ClientTimeout(total=5)
|
|
|
|
data = aiohttp.FormData()
|
|
data.add_field('files[]', image, filename='lol.jpg', content_type='image/jpg')
|
|
|
|
try:
|
|
async with aiohttp.ClientSession(
|
|
timeout=timeout,
|
|
headers={"User-Agent": "w0w"}) as session:
|
|
async with session.post(url, data=data) as resp:
|
|
blob = await resp.json()
|
|
return blob['files'][0]['url']
|
|
except Exception as ex:
|
|
raise Exception("Could not upload file to uguu.se :((")
|
|
|
|
async def get_images(self) -> Optional[List[JpegImageFile]]:
|
|
url = 'https://backend.craiyon.com/generate'
|
|
timeout_secs = 120
|
|
timeout = aiohttp.ClientTimeout(total=timeout_secs)
|
|
images = []
|
|
|
|
try:
|
|
async with aiohttp.ClientSession(
|
|
timeout=timeout,
|
|
headers=HEADERS) as session:
|
|
async with session.post(url, json={'prompt': self.term}) as resp:
|
|
if resp.status != 200:
|
|
raise Exception("POST did not return status 200")
|
|
blob = await resp.json()
|
|
if 'images' not in blob:
|
|
raise Exception("result did not return images")
|
|
images = blob['images']
|
|
except asyncio.TimeoutError:
|
|
raise Exception(f"image generation timeout exceeded ({timeout_secs} secs)")
|
|
|
|
try:
|
|
images = [b64decode(i) for i in images]
|
|
return [Image.open(BytesIO(i)) for i in images]
|
|
except Exception as ex:
|
|
raise Exception("could not decode results")
|
|
|
|
|
|
@bot.on('CLIENT_CONNECT')
|
|
async def connect(**kwargs):
|
|
bot.send('NICK', nick=IRC_NICK)
|
|
bot.send('USER', user=IRC_NICK, realname=IRC_NICK)
|
|
done, pending = await asyncio.wait(
|
|
[bot.wait("RPL_ENDOFMOTD"), bot.wait("ERR_NOMOTD")],
|
|
loop=bot.loop,
|
|
return_when=asyncio.FIRST_COMPLETED
|
|
)
|
|
|
|
for future in pending:
|
|
future.cancel()
|
|
for c in IRC_CHANNELS:
|
|
bot.send('JOIN', channel=c)
|
|
|
|
|
|
@bot.on('PING')
|
|
def keepalive(message, **kwargs):
|
|
bot.send('PONG', message=message)
|
|
|
|
|
|
@bot.on('client_disconnect')
|
|
def reconnect(**kwargs):
|
|
logging.warning("Lost IRC server connection")
|
|
time.sleep(3)
|
|
bot.loop.create_task(bot.connect())
|
|
logging.warning("Reconnecting to IRC server")
|
|
|
|
|
|
@bot.on('PRIVMSG')
|
|
async def message_received(nick, target, message, **kwargs):
|
|
if nick == IRC_NICK:
|
|
return
|
|
if target == IRC_NICK:
|
|
target = nick
|
|
if target not in IRC_CHANNELS:
|
|
return
|
|
|
|
msg = message
|
|
msg = msg.replace("!dall", " !dall")
|
|
msg = re.sub(r'!\W+', ' ', msg)
|
|
now = datetime.now()
|
|
now = time.mktime(now.timetuple())
|
|
await handle_msg(nick, msg, target, now)
|
|
|
|
|
|
async def handle_msg(nick, msg, target, now: float):
|
|
global TASK_AUTHORS, TASK_QUEUE
|
|
|
|
def _err(target, err):
|
|
return bot.send("PRIVMSG", target=target, message=err)
|
|
|
|
if "re @" in msg:
|
|
return
|
|
if "!dall" not in msg:
|
|
return
|
|
|
|
msg = msg[msg.find("!dall"):]
|
|
msg = msg[5:].strip()
|
|
|
|
if len(msg) <= 8:
|
|
return _err(target, f"longer query required.")
|
|
|
|
TASK_AUTHORS.setdefault(nick, 0)
|
|
if TASK_AUTHORS[nick] >= 3:
|
|
err = f"{nick}: you already queued 3 thingies, patient!!11"
|
|
bot.send("PRIVMSG", target=target, message=err)
|
|
return
|
|
|
|
TASK_AUTHORS[nick] += 1
|
|
task = Task(channel=target, author=nick, term=msg, dt=int(now))
|
|
await TASK_QUEUE.put(task)
|
|
bot.send("PRIVMSG", target=target, message=f"{nick}: generating...")
|
|
|
|
|
|
async def main():
|
|
global TASK_AUTHORS, TASK_QUEUE
|
|
bot.loop.create_task(bot.connect())
|
|
|
|
def _lower_author_task_count(task_author):
|
|
TASK_AUTHORS.setdefault(task_author, 1)
|
|
TASK_AUTHORS[task_author] -= 1
|
|
|
|
# main loop
|
|
while True:
|
|
random.shuffle(TASK_QUEUE._queue)
|
|
task = await TASK_QUEUE.get()
|
|
|
|
try:
|
|
images = await task.get_images()
|
|
except Exception as ex:
|
|
bot.send("PRIVMSG", target=task.channel, message=f"could not fetch images: {ex}")
|
|
_lower_author_task_count(task.author)
|
|
continue
|
|
|
|
try:
|
|
image = await task.process_images(images)
|
|
except Exception as ex:
|
|
bot.send("PRIVMSG", target=task.channel, message=f"could not process images: {ex}")
|
|
_lower_author_task_count(task.author)
|
|
continue
|
|
|
|
image_bytes = BytesIO()
|
|
image.save(image_bytes, 'JPEG')
|
|
image.save(os.path.join('data', f"{task.dt}_{task.author}.jpg"), 'JPEG')
|
|
image_bytes.seek(0)
|
|
|
|
try:
|
|
url = await task.post_image(image_bytes)
|
|
except Exception as ex:
|
|
bot.send("PRIVMSG", target=task.channel, message=str(ex))
|
|
_lower_author_task_count(task.author)
|
|
continue
|
|
|
|
_lower_author_task_count(task.author)
|
|
|
|
bot.send("PRIVMSG", target=task.channel, message=f"{url} ({task.author})")
|
|
await asyncio.sleep(10)
|
|
|
|
loop = asyncio.get_event_loop()
|
|
loop.run_until_complete(main())
|