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