diff --git a/matrixai/__init__.py b/matrixai/__init__.py index 962d727..e81e795 100644 --- a/matrixai/__init__.py +++ b/matrixai/__init__.py @@ -13,9 +13,11 @@ PASSWORD = os.environ.get("MATRIX_PASSWORD", None) ENABLE_AI_COMMAND = os.environ.get("ENABLE_AI_COMMAND", True) ENABLE_IMAGE_COMMAND = os.environ.get("ENABLE_IMAGE_COMMAND", True) +ENABLE_INSPIRE_COMMAND = os.environ.get("ENABLE_INSPIRE_COMMAND", True) AI_COMMAND_ALIASES = os.environ.get("AI_COMMAND_ALIASES", "ask, ai, gpt").split(", ") IMAGE_COMMAND_ALIASES = os.environ.get("IMAGE_COMMAND_ALIASES", "img, i").split(", ") +INSPIRE_COMMAND_ALIASES = os.environ.get("INSPIRE_COMMAND_ALIASES", "inspire, insp").split(", ") PING_URL = os.environ.get( "PING_URL", @@ -104,10 +106,6 @@ def run(): else: prompt += arg + " " - - - - async def generate_image(image_prompt, style_value, ratio_value, negative): if negative is None: negative = False @@ -148,10 +146,64 @@ def run(): async with session.get(PING_URL) as resp: await resp.read() + @bot.listener.on_message_event + async def inspire(room, message): + match = botlib.MessageMatch(room, message, bot, PREFIX) + + if match.is_not_from_this_bot() and match.prefix() and ENABLE_INSPIRE_COMMAND: + for alias in INSPIRE_COMMAND_ALIASES: + if match.command(alias): + break + else: + return + + try: + print(match.args()) + inspiration = match.args() + + async def generate_image(inspiration): + imagine = Imagine() + filename = str(uuid.uuid4()) + ".png" + try: + inspiration = Inspiration[inspiration] + except: + inspiration = Inspiration.INSPIRATION_01 + + img_data = imagine.sdinsp( + inspiration=inspiration + ) + + try: + with open(filename, mode="wb") as img_file: + img_file.write(img_data) + except Exception as e: + print( + f"An error occurred while creating the image file: {e}") + return None + + return filename + + filename = await generate_image(inspiration) + + await bot.api.send_image_message( + room_id=room.room_id, image_filepath=filename + ) + + except Exception as e: + print(e) + traceback.print_exc() + await bot.api.send_markdown_message(room.room_id, f"> {prompt}\n\n{e}") + + else: + async with aiohttp.ClientSession() as session: + async with session.get(PING_URL) as resp: + await resp.read()s + @bot.listener.on_message_event async def bot_help(room, message): styles = ", ".join([style.name for style in Style]) ratios = ", ".join([ratio.name for ratio in Ratio]) + inspirations = ", ".join([inspiration.name for inspiration in Inspiration]) bot_help_message = f""" Help Message: prefix: {PREFIX} @@ -177,6 +229,14 @@ Help Message: description: ratio of image example: ratio=RATIO_1X1 values: {ratios} + inspire: + command: {", ".join(INSPIRE_COMMAND_ALIASES)} + description: generate an image from an inspiration + options: + inspiration: + description: inspiration of image + example: INSPIRATION_01 + values: {inspirations} """ match = botlib.MessageMatch(room, message, bot, PREFIX)