Update run.py

run.py

@@ -1,39 +1,40 @@
 from flask import Flask, render_template, request, send_file, jsonify
 from imaginepy import AsyncImagine, Style, Ratio
 import os
-import aiohttp
+import asyncio
 
 app = Flask(__name__)
 
 async def generate_image_async(prompt, style, ratio):
-    async with aiohttp.ClientSession() as session:
-        imagine = AsyncImagine()
+    imagine = AsyncImagine()
 
-        try:
-            img_data = await imagine.sdprem(
-                prompt=prompt,
-                style=Style[style],
-                ratio=Ratio[ratio]
-            )
-        except Exception as e:
-            return f"An error occurred while generating the image: {e}"
+    try:
+        img_data = await imagine.sdprem(
+            prompt=prompt,
+            style=Style[style],
+            ratio=Ratio[ratio]
+        )
+    except Exception as e:
+        return f"An error occurred while generating the image: {e}"
 
-        if img_data is None:
-            return "An error occurred while generating the image."
+    if img_data is None:
+        return "An error occurred while generating the image."
 
-        img_data = await imagine.upscale(image=img_data)
+    img_data = await imagine.upscale(image=img_data)
 
-        if img_data is None:
-            return "An error occurred while upscaling the image."
+    if img_data is None:
+        return "An error occurred while upscaling the image."
 
-        try:
-            image_path = os.path.join(app.root_path, "static", "example.jpeg")
-            with open(image_path, mode="wb") as img_file:
-                img_file.write(img_data)
-        except Exception as e:
-            return f"An error occurred while writing the image to file: {e}"
+    await imagine.close()
 
-        return image_path
+    try:
+        image_path = os.path.join(app.root_path, "static", "example.jpeg")
+        with open(image_path, mode="wb") as img_file:
+            img_file.write(img_data)
+    except Exception as e:
+        return f"An error occurred while writing the image to file: {e}"
+
+    return image_path
 
 @app.route('/')
 def index():
@@ -67,5 +68,7 @@ async def api_generate_image():
     return send_file(image_path, mimetype='image/jpeg', as_attachment=True)
 
 if __name__ == "__main__":
+    loop = asyncio.get_event_loop()
+    loop.run_until_complete(main())
     app.run(host="0.0.0.0", port=7860)