Fixes metadata related to new args

This commit is contained in:
psychedelicious 2022-09-17 06:16:16 +10:00
parent d1de1e357a
commit f9feaac8c7
3 changed files with 13 additions and 14 deletions

View File

@ -147,15 +147,12 @@ def handle_request_all_images():
paths.sort(key=lambda x: os.path.getmtime(x)) paths.sort(key=lambda x: os.path.getmtime(x))
image_array = [] image_array = []
for path in paths: for path in paths:
image = Image.open(path) # image = Image.open(path)
metadata = {} all_metadata = retrieve_metadata(path)
if 'Dream' in image.info: if 'Dream' in all_metadata and not all_metadata['sd-metadata']:
try: metadata = vars(parser.parse_args(shlex.split(all_metadata['Dream'])))
metadata = vars(parser.parse_args(shlex.split(image.info['Dream']))) else:
except SystemExit: metadata = all_metadata['sd-metadata']
# TODO: Unable to parse metadata, ignore it for now,
# this can happen when metadata is missing a prompt
pass
image_array.append({'path': path, 'metadata': metadata}) image_array.append({'path': path, 'metadata': metadata})
return make_response("OK", data=image_array) return make_response("OK", data=image_array)
@ -308,7 +305,7 @@ def save_image(image, parameters, output_dir, step_index=None, postprocessing=Fa
command = parameters_to_command(parameters) command = parameters_to_command(parameters)
path = pngwriter.save_image_and_prompt_to_png(image, command, filename) path = pngwriter.save_image_and_prompt_to_png(image, command, parameters, filename)
return path return path

View File

@ -47,7 +47,8 @@ class PngWriter:
metadata stored there, as a dict metadata stored there, as a dict
''' '''
path = os.path.join(self.outdir,img_basename) path = os.path.join(self.outdir,img_basename)
return retrieve_metadata(path) all_metadata = retrieve_metadata(path)
return all_metadata['sd-metadata']
def retrieve_metadata(img_path): def retrieve_metadata(img_path):
''' '''
@ -55,6 +56,7 @@ def retrieve_metadata(img_path):
metadata stored there, as a dict metadata stored there, as a dict
''' '''
im = Image.open(img_path) im = Image.open(img_path)
md = im.text.get('sd-metadata',{}) md = im.text.get('sd-metadata', '{}')
return json.loads(md) dream_prompt = im.text.get('Dream', '')
return {'sd-metadata': json.loads(md), 'Dream': dream_prompt}

View File

@ -13,7 +13,7 @@ filenames = sys.argv[1:]
for f in filenames: for f in filenames:
try: try:
metadata = retrieve_metadata(f) metadata = retrieve_metadata(f)
print(f'{f}:\n',json.dumps(metadata, indent=4)) print(f'{f}:\n',json.dumps(metadata['sd-metadata'], indent=4))
except FileNotFoundError: except FileNotFoundError:
sys.stderr.write(f'{f} not found\n') sys.stderr.write(f'{f} not found\n')
continue continue