File size: 871 Bytes
09fee18 67f5019 7a1f675 552a266 67f5019 552a266 67f5019 62d4cc8 67f5019 09fee18 67f5019 09fee18 ff6166b 67f5019 552a266 67f5019 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | 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",
)
|