|
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, 'PNG') |
|
img_io.seek(0) |
|
return send_file(img_io, mimetype='image/png') |
|
|
|
|
|
if __name__ == '__main__': |
|
app.run('0.0.0.0', 8282, debug=False) |
