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")

@app.route(rule="/", methods=["GET"])
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",
    )