import flask
from flask import Flask, request, json, send_file, Response
import torch
from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler
from random import randrange
from io import BytesIO

repo = "Bingsu/my-korean-stable-diffusion-v1-5"
euler_ancestral_scheduler = EulerAncestralDiscreteScheduler.from_config(repo, subfolder="scheduler")
pipe = StableDiffusionPipeline.from_pretrained(
    repo, scheduler=euler_ancestral_scheduler, torch_dtype=torch.float16,
)
pipe.to("cuda")

app = Flask(__name__)

@app.post('/sd')
def generate():
    text = request.json['text']
    seed = randrange(1, 9999999999)
    generator = torch.Generator('cuda').manual_seed(seed)
    image = pipe(text, num_inference_steps=25, generator=generator).images[0]
    img_io = BytesIO()
    image.save(img_io)
    img_io.seek(0)
    return send_file(img_io, mimetype='image/png')


if __name__ == '__main__':
    app.run('0.0.0.0', 8282, debug=False)