commit 8624105c9dd179bde15c560d3dc842a0130870c9 Author: Christoph Stahl Date: Tue Jul 26 16:39:08 2022 +0200 Initial version 0.2.1 diff --git a/craiyonbot.py b/craiyonbot.py new file mode 100644 index 0000000..f5acd9f --- /dev/null +++ b/craiyonbot.py @@ -0,0 +1,140 @@ +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) diff --git a/maubot.yaml b/maubot.yaml new file mode 100644 index 0000000..03c1ca8 --- /dev/null +++ b/maubot.yaml @@ -0,0 +1,10 @@ +maubot: 0.1.0 +id: de.k-fortytwo.craiyon +version: 0.2.1 +license: AGPL-3.0-or-later +modules: + - craiyonbot +main_class: CraiyonBot +dependencies: + - python-magic>=0.4 + - Pillow>=5.1