From 7da735f847840c0d8fa511500ef0e5f331df8fd9 Mon Sep 17 00:00:00 2001 From: Gadi Cohen Date: Fri, 2 Sep 2022 16:40:41 +0300 Subject: [PATCH] Add modelInputs width, height, num_inference_steps, guidance_scale --- app.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/app.py b/app.py index 1e5aee9..7f34521 100644 --- a/app.py +++ b/app.py @@ -25,10 +25,21 @@ def inference(model_inputs:dict) -> dict: prompt = model_inputs.get('prompt', None) if prompt == None: return {'message': "No prompt provided"} + + height = model_inputs.get('height', 512); + width = model_inputs.get('width', 512); + num_inference_steps = model_inputs.get('num_inference_steps', 50); + guidance_scale = model_inputs.get('guidance_scale', 7.5); # Run the model with autocast("cuda"): - image = model(prompt)["sample"][0] + image = model( + prompt=prompt, + height=height, + width=width, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + )["sample"][0] buffered = BytesIO() image.save(buffered,format="JPEG")