From 56371153b545e3a43c3a5f206264019af361f3af Mon Sep 17 00:00:00 2001
From: DoTheSneedful <DoTheSneedful@outlook.com>
Date: Tue, 4 Oct 2022 01:07:36 -0400
Subject: [PATCH] XY plot prompt order simplify logic

---
 scripts/xy_grid.py | 22 ++++++++++++++--------
 1 file changed, 14 insertions(+), 8 deletions(-)

diff --git a/scripts/xy_grid.py b/scripts/xy_grid.py
index 5bcd3921..7def47f5 100644
--- a/scripts/xy_grid.py
+++ b/scripts/xy_grid.py
@@ -38,15 +38,21 @@ def apply_order(p, x, xs):
 
     token_order.sort(key=lambda t: t[0])
 
-    search_from_pos = 0
-    for idx, (original_pos, old_token) in enumerate(token_order):
-        # Get position of the token again as it will likely change as tokens are being replaced
-        pos = search_from_pos + p.prompt[search_from_pos:].find(old_token)
-        if original_pos >= 0:
-            # Avoid trying to replace what was just replaced by searching later in the prompt string
-            p.prompt = p.prompt[0:search_from_pos] + p.prompt[search_from_pos:].replace(old_token, x[idx], 1)
+    prompt_parts = []
 
-        search_from_pos = pos + len(x[idx])
+    # Split the prompt up, taking out the tokens
+    for _, token in token_order:
+        n = p.prompt.find(token)
+        prompt_parts.append(p.prompt[0:n])
+        p.prompt = p.prompt[n + len(token):]
+
+    # Rebuild the prompt with the tokens in the order we want
+    prompt_tmp = ""
+    for idx, part in enumerate(prompt_parts):
+        prompt_tmp += part
+        prompt_tmp += x[idx]
+    p.prompt = prompt_tmp + p.prompt
+    
 
 samplers_dict = {}
 for i, sampler in enumerate(modules.sd_samplers.samplers):