From a9b9eee182e3b4aeee9d53d4e00b2c272a3224b6 Mon Sep 17 00:00:00 2001
From: 0xMRTT <0xMRTT@proton.me>
Date: Wed, 21 Jun 2023 11:31:18 +0200
Subject: [PATCH] fix: imagine

---
 matrixai/__init__.py | 44 ++++++++++++++++++++++++++++----------------
 1 file changed, 28 insertions(+), 16 deletions(-)

diff --git a/matrixai/__init__.py b/matrixai/__init__.py
index 79505a3..f9f8ffa 100644
--- a/matrixai/__init__.py
+++ b/matrixai/__init__.py
@@ -104,24 +104,36 @@ def run():
                     else:
                         prompt += arg + " "
 
-                imagine = Imagine()
 
-                filename = str(uuid.uuid4()) + ".png"
-                # style_enum = Style[style]
-                # ratio_enum = Ratio[ratio]
+               
+                
 
-                img_data = await imagine.sdprem(
-                    prompt=prompt,
-                    #style=style_enum,
-                    #ratio=ratio_enum,
-                    negative=negative,
-                )
-                try:
-                    with open(filename, mode="wb") as img_file:
-                        img_file.write(img_data)
-                except Exception as e:
-                    print(f"An error occurred while writing the image to file: {e}")
-                    return None
+                async def generate_image(image_prompt, style_value, ratio_value, negative):
+                    if negative is None:
+                        negative = False
+                    imagine = Imagine()
+                    filename = str(uuid.uuid4()) + ".png"
+                    style_enum = Style[style_value]
+                    ratio_enum = Ratio[ratio_value]
+                    img_data = await imagine.sdprem(
+                        prompt=image_prompt,
+                        style=style_enum,
+                        ratio=ratio_enum,
+                        negative=negative
+                    )
+
+                    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
+
+                    await imagine.close()
+                    return img_file
+                
+                filename = await generate_image(prompt, style, ratio, negative)
 
                 await bot.api.send_image_message(
                     room_id=room.room_id, image_filepath=filename