Imagine / run.py
imseldrith's picture
Update run.py
318f8d9
raw
history blame
2.04 kB
from flask import Flask, render_template, request, send_file, jsonify
from imaginepy import AsyncImagine, Style, Ratio
import os
import asyncio
app = Flask(__name__)
async def generate_image_async(prompt, style, ratio):
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}"
if img_data is None:
return "An error occurred while generating the image."
img_data = await imagine.upscale(image=img_data)
if img_data is None:
return "An error occurred while upscaling the image."
await imagine.close()
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():
return render_template('index.html')
@app.route('/generate', methods=['POST'])
async def generate_image():
prompt = request.form['prompt']
style = request.form['style']
ratio = request.form['ratio']
image_path = await generate_image_async(prompt, style, ratio)
if isinstance(image_path, str):
return image_path
return render_template('output.html')
@app.route('/api/generate', methods=['POST'])
async def api_generate_image():
data = request.get_json()
prompt = data['prompt']
style = data['style']
ratio = data['ratio']
image_path = await generate_image_async(prompt, style, ratio)
if isinstance(image_path, str):
return jsonify({'error': image_path}), 500
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)