apextoaster
/
onnx-web
1
0
Fork 0
Code Actions

add a basic lock

This commit is contained in:
Sean Sube 2023-01-04 19:51:09 -06:00
parent 3cb527c2b8
commit 4c9a0b1ead
1 changed files with 12 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
api/serve.py
View File

@ -14,6 +14,9 @@ max_cfg = 8
model_path = environ.get('ONNX_WEB_MODEL_PATH', "../../stable_diffusion_onnx") model_path = environ.get('ONNX_WEB_MODEL_PATH', "../../stable_diffusion_onnx")
output_path = environ.get('ONNX_WEB_OUTPUT_PATH', "../../web_output") output_path = environ.get('ONNX_WEB_OUTPUT_PATH', "../../web_output")
# queue
image_queue = set()
def setup(): def setup():
if not path.exists(model_path): if not path.exists(model_path):
raise RuntimeError('model path must exist') raise RuntimeError('model path must exist')
@ -32,17 +35,25 @@ def hello():
@app.route('/txt2img') @app.route('/txt2img')
def txt2img(): def txt2img():
if len(image_queue) > 0:
return 'Queue full: %s' % (image_queue)
origin = request.origin
prompt = request.args.get('prompt', empty_prompt) prompt = request.args.get('prompt', empty_prompt)
height = request.args.get('height', max_height) height = request.args.get('height', max_height)
width = request.args.get('width', max_width) width = request.args.get('width', max_width)
steps = int(request.args.get('steps', max_steps)) steps = int(request.args.get('steps', max_steps))
cfg = int(request.args.get('cfg', max_cfg)) cfg = int(request.args.get('cfg', max_cfg))
print("txt2img: %s/%s, %sx%s, %s" % (cfg, steps, width, height, prompt)) print("txt2img from %s: %s/%s, %sx%s, %s" % (origin, cfg, steps, width, height, prompt))
image_queue.add(origin)
image = pipe(prompt, height, width, num_inference_steps=steps, guidance_scale=cfg).images[0] image = pipe(prompt, height, width, num_inference_steps=steps, guidance_scale=cfg).images[0]
# image.save("astronaut_rides_horse.png") # image.save("astronaut_rides_horse.png")
img_io = BytesIO() img_io = BytesIO()
image.save(img_io, 'PNG', quality=100) image.save(img_io, 'PNG', quality=100)
img_io.seek(0) img_io.seek(0)
image_queue.remove(origin)
return send_file(img_io, mimetype='image/png') return send_file(img_io, mimetype='image/png')