Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix long-standing issue with metadata retrieval

Browse Source 
The Args object would crap out when trying to retrieve metadata from
an image file that did not contain InvokeAI-generated metadata, such
as a JPG. This corrects that and returns dummy values (seed of zero,
prompt of '') to avoid downstream breakage.
This commit is contained in:
Lincoln Stein 2022-10-27 22:43:34 -04:00
parent 7b46d5f823
commit 362b234cd1
1 changed files with 23 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
ldm/invoke/args.py
View File

@ -175,7 +175,7 @@ class Args(object):
switches = cmd_string switches = cmd_string
# handle the case in which the prompt is enclosed by quotes # handle the case in which the prompt is enclosed by quotes
elif cmd_string.startswith('"'): elif cmd_string.startswith('"'):
a = shlex.split(cmd_string) a = shlex.split(cmd_string,comments=True)
prompt = a[0] prompt = a[0]
switches = shlex.join(a[1:]) switches = shlex.join(a[1:])
else: else:
@ -192,7 +192,7 @@ class Args(object):
prompt = cmd_string prompt = cmd_string
switches = '' switches = ''
try: try:
self._cmd_switches = self._cmd_parser.parse_args(shlex.split(switches)) self._cmd_switches = self._cmd_parser.parse_args(shlex.split(switches,comments=True))
setattr(self._cmd_switches,'prompt',prompt) setattr(self._cmd_switches,'prompt',prompt)
return self._cmd_switches return self._cmd_switches
except: except:
@ -932,7 +932,7 @@ def metadata_dumps(opt,
# remove any image keys not mentioned in RFC #266 # remove any image keys not mentioned in RFC #266
rfc266_img_fields = ['type','postprocessing','sampler','prompt','seed','variations','steps', rfc266_img_fields = ['type','postprocessing','sampler','prompt','seed','variations','steps',
'cfg_scale','threshold','perlin','fnformat', 'step_number','width','height','extra','strength', 'cfg_scale','threshold','perlin','step_number','width','height','extra','strength',
'init_img','init_mask','facetool','facetool_strength','upscale'] 'init_img','init_mask','facetool','facetool_strength','upscale']
rfc_dict ={} rfc_dict ={}
@ -983,6 +983,23 @@ def metadata_dumps(opt,
return metadata return metadata
@functools.lru_cache(maxsize=50)
def args_from_png(png_file_path) -> list[Args]:
'''
Given the path to a PNG file created by invoke.py,
retrieves a list of Args objects containing the image
data.
'''
try:
meta = ldm.invoke.pngwriter.retrieve_metadata(png_file_path)
except AttributeError:
return [legacy_metadata_load({},png_file_path)]
try:
return metadata_loads(meta)
except:
return [legacy_metadata_load(meta,png_file_path)]
@functools.lru_cache(maxsize=50) @functools.lru_cache(maxsize=50)
def metadata_from_png(png_file_path) -> Args: def metadata_from_png(png_file_path) -> Args:
''' '''
@ -990,11 +1007,8 @@ def metadata_from_png(png_file_path) -> Args:
an Args object containing the image metadata. Note that this an Args object containing the image metadata. Note that this
returns a single Args object, not multiple. returns a single Args object, not multiple.
''' '''
meta = ldm.invoke.pngwriter.retrieve_metadata(png_file_path) args_list = args_from_png(png_file_path)
if 'sd-metadata' in meta and len(meta['sd-metadata'])>0 : return args_list[0]
return metadata_loads(meta)[0]
else:
return legacy_metadata_load(meta,png_file_path)
def dream_cmd_from_png(png_file_path): def dream_cmd_from_png(png_file_path):
opt = metadata_from_png(png_file_path) opt = metadata_from_png(png_file_path)
@ -1009,7 +1023,7 @@ def metadata_loads(metadata) -> list:
''' '''
results = [] results = []
try: try:
if 'grid' in metadata['sd-metadata']: if 'images' in metadata['sd-metadata']:
images = metadata['sd-metadata']['images'] images = metadata['sd-metadata']['images']
else: else:
images = [metadata['sd-metadata']['image']] images = [metadata['sd-metadata']['image']]