From 107bbfa2919ed084cc79cc5f62bc7cfd93d27770 Mon Sep 17 00:00:00 2001 From: Robert Vining Date: Tue, 1 Nov 2022 18:01:20 +1000 Subject: [PATCH] Added new function text2img_api() - sends prompt to locally hosted stable-diffusion-webui API and retrieves image result. info from SD API is embedded in png file as it is saved. image files are named with the current date at generation and saved in /stories/art/ Modified UI_2_generate_image() to call text2img_api() instead of text2img_horde Signed-off-by: Robert Vining --- aiserver.py | 58 ++++++++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 53 insertions(+), 5 deletions(-) diff --git a/aiserver.py b/aiserver.py index 422cbed9..5abec854 100644 --- a/aiserver.py +++ b/aiserver.py @@ -81,7 +81,7 @@ import transformers.generation_utils # Text2img import base64 -from PIL import Image, ImageFont, ImageDraw, ImageFilter, ImageOps +from PIL import Image, ImageFont, ImageDraw, ImageFilter, ImageOps, PngImagePlugin from io import BytesIO global tpu_mtj_backend @@ -9098,7 +9098,7 @@ def UI_2_generate_image(data): #If we don't have a GPU, use horde if we're allowed to start_time = time.time() if ((not koboldai_vars.hascuda or not os.path.exists("models/stable-diffusion-v1-4")) and koboldai_vars.img_gen_priority != 0) or koboldai_vars.img_gen_priority == 3: - b64_data = text2img_horde(", ".join(keys), art_guide = art_guide) + b64_data = text2img_api(", ".join(keys), art_guide = art_guide) else: import psutil #We aren't being forced to use horde, so now let's figure out if we should use local @@ -9108,9 +9108,9 @@ def UI_2_generate_image(data): elif torch.cuda.get_device_properties(0).total_memory > 6000000000 and koboldai_vars.img_gen_priority <= 1: #We could do it locally by swapping the model out print("Could do local or online") - b64_data = text2img_horde(", ".join(keys), art_guide = art_guide) + b64_data = text2img_api(", ".join(keys), art_guide = art_guide) elif koboldai_vars.img_gen_priority != 0: - b64_data = text2img_horde(", ".join(keys), art_guide = art_guide) + b64_data = text2img_api(", ".join(keys), art_guide = art_guide) logger.debug("Time to Generate Image {}".format(time.time()-start_time)) koboldai_vars.picture = b64_data koboldai_vars.picture_prompt = ", ".join(keys) @@ -9136,7 +9136,7 @@ def text2img_local(prompt, art_guide="", filename="new.png"): def get_image(pipe, prompt, num_inference_steps): from torch import autocast with autocast("cuda"): - return pipe(prompt, num_inference_steps=num_inference_steps)["sample"][0] + return pipe(prompt, num_inference_steps=num_inference_steps).images[0] image = tpool.execute(get_image, pipe, prompt, num_inference_steps=35) buffered = BytesIO() image.save(buffered, format="JPEG") @@ -9195,6 +9195,54 @@ def text2img_horde(prompt, koboldai_vars.generating_image = False logger.error(submit_req.text) +#@logger.catch +def text2img_api(prompt, + art_guide = 'fantasy illustration, artstation, by Hugin Miyama by Taiki Kawakami, cinematic lighting', + filename = "story_art.png"): + logger.debug("Generating Image using external API") + koboldai_vars.generating_image = True + final_imgen_params = { + "prompt": "{}, {}".format(prompt, art_guide), + "n": 1, + "width": 512, + "height": 512, + "steps": 50, + "cfg_scale": 14, + "negative_prompt": " lowres, bad anatomy, bad hands out of frame, two heads, totem pole, several faces, extra fingers, mutated hands, (poorly drawn hands:1.21), (poorly drawn face:1.21), (mutation:1.331), (deformed:1.331), (ugly:1.21), blurry, (bad anatomy:1.21), (bad proportions:1.331), (extra limbs:1.21), glitchy, ((clip through table)), adherent bodies, slimy bodies, (badly visible legs), captions, words", + } + + final_submit_dict = { + "prompt": "{}, {}".format(prompt, art_guide), + "params": final_imgen_params, + } + apiaddress = 'http://127.0.0.1:7860/sdapi/v1/txt2img' + payload_json = json.dumps(final_imgen_params) + logger.debug(final_submit_dict) + submit_req = requests.post(url=f'{apiaddress}', data=payload_json).json() + if submit_req: + results = submit_req + for i in results['images']: + final_src_img = Image.open(BytesIO(base64.b64decode(i.split(",",1)[1]))) + buffer = BytesIO() + final_src_img.save(buffer, format="Webp", quality=95) + b64img = base64.b64encode(buffer.getvalue()).decode("utf8") + base64_bytes = b64img.encode('utf-8') + img_bytes = base64.b64decode(base64_bytes) + img = Image.open(BytesIO(img_bytes)) + #dtnow = datetime.datetime.now() + #dt_string = dtnow().strftime("%H%M%S%d%m%Y") + dt_string = datetime.datetime.now().strftime("%H%M%S%d%m%Y") + final_filename = "stories/art/{}_{}".format(dt_string,filename) + pnginfo = PngImagePlugin.PngInfo() + pnginfo.add_text("parameters", str(results['info'])) + img.save(final_filename, pnginfo=pnginfo) + logger.debug("Saved Image") + koboldai_vars.generating_image = False + return(b64img) + else: + koboldai_vars.generating_image = False + logger.error(submit_req.text) + #@logger.catch def get_items_locations_from_text(text): # load model and tokenizer