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:")