rixty/KoboldAI-Client
1
1
Fork 0
mirror of https://github.com/KoboldAI/KoboldAI-Client.git synced 2025-06-05 21:59:24 +02:00
Code Issues Projects Releases Wiki Activity

Added EXIF info on stable diffusion images

Browse Source
This commit is contained in:
ebolam
2022-12-14 11:15:44 -05:00
parent 50f9508ee7
commit b25cc97c40
1 changed files with 26 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
aiserver.py
View File

@@ -9604,25 +9604,7 @@ def UI_2_generate_image_from_story(data):
else: else:
action = koboldai_vars.prompt action = koboldai_vars.prompt
action_id = -1 action_id = -1
#Get matching world info entries
keys = []
for wi in koboldai_vars.worldinfo_v2:
for key in wi['key']:
if key in action:
#Check to make sure secondary keys are present if needed
if len(wi['keysecondary']) > 0:
for keysecondary in wi['keysecondary']:
if keysecondary in action:
keys.append(key)
break
break
else:
keys.append(key)
break
#If we have > 4 keys, use those otherwise use sumarization
if len(keys) < 4:
start_time = time.time() start_time = time.time()
if os.path.exists("models/{}".format(args.summarizer_model.replace('/', '_'))): if os.path.exists("models/{}".format(args.summarizer_model.replace('/', '_'))):
koboldai_vars.summary_tokenizer = AutoTokenizer.from_pretrained("models/{}".format(args.summarizer_model.replace('/', '_')), cache_dir="cache") koboldai_vars.summary_tokenizer = AutoTokenizer.from_pretrained("models/{}".format(args.summarizer_model.replace('/', '_')), cache_dir="cache")
@@ -9716,16 +9698,20 @@ def generate_story_image(
if not image: if not image:
return return
exif = image.getexif()
exif[0x9286] = prompt
exif[0x927C] = generation_type if generation_type != "" else "Stable Diffusion from KoboldAI"
if os.path.exists(koboldai_vars.save_paths.generated_images): if os.path.exists(koboldai_vars.save_paths.generated_images):
# Only save image if this is a saved story # Only save image if this is a saved story
file_name = f"{file_prefix}_{int(time.time())}.png" file_name = f"{file_prefix}_{int(time.time())}.jpg"
image.save(os.path.join(koboldai_vars.save_paths.generated_images, file_name)) image.save(os.path.join(koboldai_vars.save_paths.generated_images, file_name), format="JPEG", exif=exif)
log_image_generation(prompt, display_prompt, file_name, generation_type, log_data) log_image_generation(prompt, display_prompt, file_name, generation_type, log_data)
logger.debug("Time to Generate Image {}".format(time.time()-start_time)) logger.debug("Time to Generate Image {}".format(time.time()-start_time))
buffer = BytesIO() buffer = BytesIO()
image.save(buffer, format="JPEG") image.save(buffer, format="JPEG", exif=exif)
b64_data = base64.b64encode(buffer.getvalue()).decode("ascii") b64_data = base64.b64encode(buffer.getvalue()).decode("ascii")
koboldai_vars.picture = b64_data koboldai_vars.picture = b64_data