Update test.py

test.py

@@ -1,7 +1,7 @@
+import asyncio
 from flask import Flask, render_template, request, send_file, jsonify
 from imaginepy import AsyncImagine, Style, Ratio
 import os
-import asyncio
 
 app = Flask(__name__)
 
@@ -15,33 +15,34 @@ async def generate_image():
     style = request.form['style']
     ratio = request.form['ratio']
 
-    imagine = AsyncImagine()
-
     try:
+        imagine = AsyncImagine()
+
         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}", 500
 
-    if img_data is None:
-        return "An error occurred while generating the image.", 500
+        if img_data is None:
+            return "An error occurred while generating the image.", 500
 
-    img_data = 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.", 500
-    
-    imagine.close()
-    try:
-        with open("static/example.jpeg", mode="wb") as img_file:
+        if img_data is None:
+            return "An error occurred while upscaling the image.", 500
+
+        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)
+
+        return render_template('output.html')
+
     except Exception as e:
-        return f"An error occurred while writing the image to file: {e}", 500
+        return f"An error occurred: {e}", 500
 
-    return render_template('output.html')
+    finally:
+        await imagine.close()
 
 @app.route('/api/generate', methods=['POST'])
 async def api_generate_image():
@@ -50,34 +51,34 @@ async def api_generate_image():
     style = data['style']
     ratio = data['ratio']
 
-    imagine = AsyncImagine()
-    
     try:
+        imagine = AsyncImagine()
+
         img_data = await imagine.sdprem(
             prompt=prompt,
             style=Style[style],
             ratio=Ratio[ratio]
         )
-    except Exception as e:
-        return jsonify({'error': f"An error occurred while generating the image: {e}"}), 500
 
-    if img_data is None:
-        return jsonify({'error': "An error occurred while generating the image."}), 500
+        if img_data is None:
+            return jsonify({'error': "An error occurred while generating the image."}), 500
 
-    img_data = imagine.upscale(image=img_data)
+        img_data = await imagine.upscale(image=img_data)
 
-    if img_data is None:
-        return jsonify({'error': "An error occurred while upscaling the image."}), 500
+        if img_data is None:
+            return jsonify({'error': "An error occurred while upscaling the image."}), 500
 
-    imagine.close()
-    try:
         image_path = os.path.join(app.root_path, "generated.jpeg")
         with open(image_path, mode="wb") as img_file:
             img_file.write(img_data)
+
+        return send_file(image_path, mimetype='image/jpeg', as_attachment=True)
+
     except Exception as e:
-        return jsonify({'error': f"An error occurred while writing the image to file: {e}"}), 500
+        return jsonify({'error': f"An error occurred: {e}"}), 500
 
-    return send_file(image_path, mimetype='image/jpeg', as_attachment=True)
+    finally:
+        await imagine.close()
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)