Update run.py

run.py

@@ -6,33 +6,34 @@ import aiohttp
 app = Flask(__name__)
 
 async def generate_image_async(prompt, style, ratio):
-    imagine = AsyncImagine()
+    async with aiohttp.ClientSession() as session:
+        imagine = AsyncImagine(session=session)
 
-    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}"
+        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
+        return image_path
 
 @app.route('/')
 def index():
@@ -67,3 +68,4 @@ async def api_generate_image():
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)
+