File size: 961 Bytes
64e6646 bf988c9 70b0399 64e6646 bf988c9 64e6646 bf988c9 64e6646 bf988c9 70b0399 81bf2a7 cb1ff21 64e6646 cc8f8af |
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 |
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) |