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