| from diffusers import AutoPipelineForText2Image | |
| from flask import ( | |
| Flask, | |
| jsonify, | |
| request, | |
| send_file, | |
| ) | |
| from io import BytesIO | |
| from torch import cuda | |
| from urllib.parse import unquote | |
| app = Flask("ImageAPI") | |
| def main(): | |
| prompt = request.args.get("prompt") | |
| if not prompt: | |
| return jsonify({ | |
| "status": 200, | |
| "message": "API is Alive!", | |
| }) | |
| pipeline = AutoPipelineForText2Image.from_pretrained( | |
| pretrained_model_or_path="Lykon/DreamShaper", | |
| ) | |
| if cuda.is_available(): | |
| pipeline.to("cuda") | |
| else: | |
| pipeline.to("cpu") | |
| image = pipeline( | |
| unquote(string=prompt), | |
| height=1024, | |
| width=1024, | |
| ).images[0] | |
| return send_file( | |
| path_or_file=BytesIO(image), | |
| mimetype="image/png", | |
| ) | |