From 4c9a0b1ead26d1448a27cf1c0227d11cb468745b Mon Sep 17 00:00:00 2001 From: Sean Sube Date: Wed, 4 Jan 2023 19:51:09 -0600 Subject: [PATCH] add a basic lock --- api/serve.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/api/serve.py b/api/serve.py index 66426c33..7d1fa08d 100644 --- a/api/serve.py +++ b/api/serve.py @@ -14,6 +14,9 @@ max_cfg = 8 model_path = environ.get('ONNX_WEB_MODEL_PATH', "../../stable_diffusion_onnx") output_path = environ.get('ONNX_WEB_OUTPUT_PATH', "../../web_output") +# queue +image_queue = set() + def setup(): if not path.exists(model_path): raise RuntimeError('model path must exist') @@ -32,17 +35,25 @@ def hello(): @app.route('/txt2img') def txt2img(): + if len(image_queue) > 0: + return 'Queue full: %s' % (image_queue) + + origin = request.origin prompt = request.args.get('prompt', empty_prompt) height = request.args.get('height', max_height) width = request.args.get('width', max_width) steps = int(request.args.get('steps', max_steps)) 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.save("astronaut_rides_horse.png") img_io = BytesIO() image.save(img_io, 'PNG', quality=100) img_io.seek(0) + + image_queue.remove(origin) return send_file(img_io, mimetype='image/png') \ No newline at end of file