Initial version 0.2.1

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 <text>
+
+    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}<br>'}
+                    )
+
+        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)

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