from base64 import b64decode from dataclasses import dataclass from io import BytesIO from typing import Tuple from maubot import Plugin, MessageEvent from maubot.handlers import command from mautrix.types import ImageInfo, EventType, ReactionEvent from PIL import Image import magic EMOJI_REGEX = r"^[\U00000031-\U00000039]\U0000FE0F\U000020E3" EMOJI_NUMBERS = ["1️⃣", "2️⃣", "3️⃣", "4️⃣", "5️⃣", "6️⃣", "7️⃣", "8️⃣", "9️⃣"] @dataclass class IMMeta: """ Lightweight wrapper for meta data for the images. Kept in memory """ uri: str name: str mime: str size: int height: int width: int class UploadableImage: """ Handles the internals for images """ def __init__(self, name, image_bytes=None, pil=None): if image_bytes is None: with BytesIO() as output: pil.save(output, format="PNG") self.image_bytes = output.getvalue() else: self.image_bytes = image_bytes self.name = name self.mime = magic.from_buffer(self.image_bytes, mime=True) self.pil = Image.open(BytesIO(image_bytes)) if pil is None else pil self.size = len(self.image_bytes) self.uri = None async def upload(self, client: Plugin) -> str: """ uploads the contained image to the matrix network. Sets it's own uri """ self.uri = await client.upload_media(self.image_bytes, self.mime, filename=self.name ) return self.uri def meta(self) -> IMMeta: """ Returns a lightweight version of the image (without the pixeldata), to be stored in memory. """ return IMMeta(self.uri, self.name, self.mime, len(self.image_bytes), self.pil.height, self.pil.width) class CraiyonBot(Plugin): """ A Maubot plugin. It registers the command !craiyon and then forwards all images in a 3 by 3 image matrix to the user. Also gives the choice to select single images by using the emoji reactions The mxc references are only stored at runtime and forgotten, once the bot restarts. """ images = {} def emoji_to_number(self, emoji: str) -> int: """ Helper function to convert the emojis to their index """ return EMOJI_NUMBERS.index(emoji) @command.passive(regex=EMOJI_REGEX, field=lambda evt: evt.content.relates_to.key, event_type=EventType.REACTION, msgtypes=None) async def get_image(self, evt: ReactionEvent, _: Tuple[str]) -> None: """ Handler when clicked a reaction. Returns the chosen image """ msg = evt.content.relates_to.event_id if msg in self.images: index = self.emoji_to_number(evt.content.relates_to.key) image = self.images[msg][index] await self.client.send_image(evt.room_id, url=image.uri, file_name=image.name, info=ImageInfo(mimetype=image.mime, size=image.size, width=image.width, height=image.height)) @command.new() @command.argument("prompt", pass_raw=True, required=True) async def craiyon(self, evt: MessageEvent, prompt: str) -> None: """ Forwards the request to craiyon and returns the images craiyon created """ await evt.react("🤖") response = await self.http.post( 'https://backend.craiyon.com/generate', json={'prompt': f'{prompt}
'} ) images = [ (n, UploadableImage(f"{prompt}_{n}.jpg", b64decode(image))) for n, image in enumerate((await response.json())['images']) ] images_canvas = Image.new(mode="RGB", size=(256*3, 256*3)) images_data = [] for index, image in images: pos_x, pos_y = index % 3, index // 3 images_canvas.paste( image.pil, (pos_x * 256, pos_y * 256) ) await image.upload(self.client) images_data.append(image.meta()) uploadable_3x3 = UploadableImage( f"{prompt}_3x3.png", pil=images_canvas) await uploadable_3x3.upload(self.client) msg = await self.client.send_image(evt.room_id, url=uploadable_3x3.uri, file_name=uploadable_3x3.name, info=ImageInfo( mimetype=uploadable_3x3.mime, size=uploadable_3x3.size, width=uploadable_3x3.pil.width, height=uploadable_3x3.pil.height) ) self.images[msg] = images_data for emoji in EMOJI_NUMBERS: await evt.client.react(evt.room_id, msg, emoji)