Merge pull request #259 from hinablue/main

Rewrite split_exif_data for more complex infotext

scripts/wib/wib_db.py

@@ -2,6 +2,7 @@ import hashlib
 import json
 import os
 import sqlite3
+import re
 from shutil import copy2
 from modules import scripts, shared
 from tempfile import gettempdir
@@ -185,6 +186,57 @@ def split_exif_data(info):
     negative_prompt = "0"
     key_values = "0: 0"
     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":
         info_list = info.split("\n")
         prompt = ""
@@ -202,20 +254,18 @@ def split_exif_data(info):
                     # multiline prompts
                     prompt = f"{prompt}\n{info_item}"
     if key_values != "":
-        key_value = ""
-        quote_open = False
-        for char in key_values + ",":
-            key_value += char
-            if char == '"':
-                quote_open = not quote_open
-            if char == "," and not quote_open:
-                try:
-                    k, v = key_value.strip(" ,").split(": ")
-                except ValueError:
-                    k = key_value.strip(" ,").split(": ")[0]
-                    v = ""
-                key_value_pairs.append((k, v))
-                key_value = ""
+        pattern = r'(\w+(?:\s+\w+)*?):\s*((?:"[^"]*"|[^,])+)(?:,\s*|$)'
+        matches = re.findall(pattern, key_values)
+        result = {key.strip(): value.strip() for key, value in matches}
+
+        # Save 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, key)
+
+            key_value_pairs.append((key, value))
 
     return prompt, negative_prompt, key_value_pairs