Rewrite split_exif_data for more complex infotext

pull/259/head
Hina Chen 2024-08-15 12:23:02 +08:00
parent 19b4b7fc8f
commit 3829185fb7
1 changed files with 94 additions and 44 deletions

View File

@ -2,6 +2,7 @@ import hashlib
import json import json
import os import os
import sqlite3 import sqlite3
import re
from shutil import copy2 from shutil import copy2
from modules import scripts, shared from modules import scripts, shared
from tempfile import gettempdir from tempfile import gettempdir
@ -185,6 +186,57 @@ def split_exif_data(info):
negative_prompt = "0" negative_prompt = "0"
key_values = "0: 0" key_values = "0: 0"
key_value_pairs = [] key_value_pairs = []
def parse_value_pairs(kv_str, key_prefix=''):
# Regular expression pattern to match key-value pairs, including multiline prompts
pattern = r'((?:\w+ )?(?:Prompt|Negative Prompt)|[^:]+):\s*((?:[^,]+(?:,(?![^:]+:))?)+)'
# Find all matches
matches = re.findall(pattern, kv_str, re.IGNORECASE | re.DOTALL)
result = {}
current_prompt = None
def process_prompt(key, value, current_prompt):
if current_prompt is None:
result[key] = value
current_prompt = key
else:
pk_values = [v.strip() for v in key.split(',') if v.strip()]
result[current_prompt] += f",{','.join(pk_values[:-1])}"
current_prompt = pk_values[-1]
result[current_prompt] = ','.join([v.strip() for v in value.split(',') if v.strip()])
return current_prompt
def process_regular_key(key, value, current_prompt):
values = [v.strip() for v in value.split(',') if v.strip()]
if current_prompt is not None:
pk_values = [v.strip() for v in key.split(',') if v.strip()]
result[current_prompt] += f",{','.join(pk_values[:-1])}"
current_prompt = None
key = pk_values[-1]
result[key] = values[0] if len(values) == 1 else ','.join(values)
return current_prompt
for key, value in matches:
key = key.strip(' ,')
value = value.strip()
if "prompt" in key.lower() or "prompt" in value.lower():
current_prompt = process_prompt(key, value, current_prompt)
else:
current_prompt = process_regular_key(key, value, current_prompt)
# Print the resulting key-value pairs
for key, value in result.items():
value = value.strip(' ,')
if value.startswith('"') and value.endswith('"'):
value = value[1:-1]
parse_value_pairs(value, f"{key_prefix} - {key}" if key_prefix != '' else key)
key_value_pairs.append((f"{key_prefix} - {key}" if key_prefix != '' else key, value))
if info != "0": if info != "0":
info_list = info.split("\n") info_list = info.split("\n")
prompt = "" prompt = ""
@ -202,20 +254,18 @@ def split_exif_data(info):
# multiline prompts # multiline prompts
prompt = f"{prompt}\n{info_item}" prompt = f"{prompt}\n{info_item}"
if key_values != "": if key_values != "":
key_value = "" pattern = r'(\w+(?:\s+\w+)*?):\s*((?:"[^"]*"|[^,])+)(?:,\s*|$)'
quote_open = False matches = re.findall(pattern, key_values)
for char in key_values + ",": result = {key.strip(): value.strip() for key, value in matches}
key_value += char
if char == '"': # Save resulting key-value pairs
quote_open = not quote_open for key, value in result.items():
if char == "," and not quote_open: value = value.strip(' ,')
try: if value.startswith('"') and value.endswith('"'):
k, v = key_value.strip(" ,").split(": ") value = value[1:-1]
except ValueError: parse_value_pairs(value, key)
k = key_value.strip(" ,").split(": ")[0]
v = "" key_value_pairs.append((key, value))
key_value_pairs.append((k, v))
key_value = ""
return prompt, negative_prompt, key_value_pairs return prompt, negative_prompt, key_value_pairs