fix prompt parser, image save

This commit is contained in:
Vladimir Mandic
2024-03-21 12:11:47 -04:00
parent 1a27871c70
commit 8a869d7e27
4 changed files with 49 additions and 32 deletions
+6 -1
View File
@@ -192,6 +192,8 @@ def parse_generation_parameters(infotext):
debug(f'Parse infotext: {infotext}')
re_param = re.compile(r'\s*([\w ]+):\s*("(?:\\"[^,]|\\"|\\|[^\"])+"|[^,]*)(?:,|$)') # multi-word: value
re_size = re.compile(r"^(\d+)x(\d+)$") # int x int
basic_params = ['steps', 'seed', 'width', 'height', 'sampler', 'size', 'cfg scale'] # first param is one of those
sanitized = infotext.replace('prompt:', 'Prompt:').replace('negative prompt:', 'Negative prompt:').replace('Negative Prompt', 'Negative prompt') # cleanup everything in brackets so re_params can work
sanitized = re.sub(r'<[^>]*>', lambda match: ' ' * len(match.group()), sanitized)
sanitized = re.sub(r'\([^)]*\)', lambda match: ' ' * len(match.group()), sanitized)
@@ -200,7 +202,10 @@ def parse_generation_parameters(infotext):
params = dict(re_param.findall(sanitized))
debug(f"Parse params: {params}")
params = { k.strip():params[k].strip() for k in params if k.lower() not in ['hashes', 'lora', 'embeddings', 'prompt', 'negative prompt']} # remove some keys
first_param = next(iter(params)) if params else None
first_param, first_param_idx = next((s, i) for i, s in enumerate(params) if any(x in s.lower() for x in basic_params))
if first_param_idx > 0:
for _i in range(first_param_idx):
params.pop(next(iter(params)))
params_idx = sanitized.find(f'{first_param}:') if first_param else -1
negative_idx = infotext.find("Negative prompt:")
+9 -6
View File
@@ -548,6 +548,7 @@ def atomically_save_image():
image = set_watermark(image, shared.opts.image_watermark)
size = os.path.getsize(fn) if os.path.exists(fn) else 0
shared.log.info(f'Saving: image="{fn}" type={image_format} resolution={image.width}x{image.height} size={size}')
exifinfo = (exifinfo or "") if shared.opts.image_metadata else ""
# additional metadata saved in files
if shared.opts.save_txt and len(exifinfo) > 0:
try:
@@ -557,7 +558,6 @@ def atomically_save_image():
except Exception as e:
shared.log.warning(f'Saving failed: description={filename_txt} {e}')
# actual save
exifinfo = (exifinfo or "") if shared.opts.image_metadata else ""
if image_format == 'PNG':
pnginfo_data = PngImagePlugin.PngInfo()
for k, v in params.pnginfo.items():
@@ -569,19 +569,22 @@ def atomically_save_image():
image = image.convert("RGB")
elif image.mode == 'I;16':
image = image.point(lambda p: p * 0.0038910505836576).convert("L")
exif_bytes = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo, encoding="unicode") } })
save_args = { 'optimize': True, 'quality': shared.opts.jpeg_quality, 'exif': exif_bytes if shared.opts.image_metadata else None }
save_args = { 'optimize': True, 'quality': shared.opts.jpeg_quality }
if shared.opts.image_metadata:
save_args['exif'] = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo, encoding="unicode") } })
elif image_format == 'WEBP':
if image.mode == 'I;16':
image = image.point(lambda p: p * 0.0038910505836576).convert("RGB")
exif_bytes = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo, encoding="unicode") } })
save_args = { 'optimize': True, 'quality': shared.opts.jpeg_quality, 'exif': exif_bytes if shared.opts.image_metadata else None, 'lossless': shared.opts.webp_lossless }
save_args = { 'optimize': True, 'quality': shared.opts.jpeg_quality, 'lossless': shared.opts.webp_lossless }
if shared.opts.image_metadata:
save_args['exif'] = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo, encoding="unicode") } })
else:
save_args = { 'quality': shared.opts.jpeg_quality }
try:
image.save(fn, format=image_format, **save_args)
except Exception as e:
shared.log.error(f'Saving failed: file="{fn}" format={image_format} {e}')
shared.log.error(f'Saving failed: file="{fn}" format={image_format} args={save_args} {e}')
errors.display(e, 'Image save')
if shared.opts.save_log_fn != '' and len(exifinfo) > 0:
fn = os.path.join(paths.data_path, shared.opts.save_log_fn)
if not fn.endswith('.json'):