From 5fbed65236d539b657eaf3b88138388cd8b66544 Mon Sep 17 00:00:00 2001
From: cryzed <cryzed@googlemail.com>
Date: Sun, 11 Sep 2022 16:35:12 +0200
Subject: Add support for saving styles with negative prompts

---
 modules/styles.py | 81 +++++++++++++++++++++++++++++++++++--------------------
 1 file changed, 52 insertions(+), 29 deletions(-)

(limited to 'modules/styles.py')

diff --git a/modules/styles.py b/modules/styles.py
index 85b78fad..25ec504b 100644
--- a/modules/styles.py
+++ b/modules/styles.py
@@ -1,44 +1,67 @@
+# We need this so Python doesn't complain about the unknown StableDiffusionProcessing-typehint at runtime
+from __future__ import annotations
+
 import csv
+import os
 import os.path
-from collections import namedtuple
+import typing
+import collections.abc as abc
+import tempfile
+import shutil
 
-PromptStyle = namedtuple("PromptStyle", ["name", "text"])
+if typing.TYPE_CHECKING:
+    # Only import this when code is being type-checked, it doesn't have any effect at runtime
+    from .processing import StableDiffusionProcessing
 
 
-def load_styles(filename):
-    res = {"None": PromptStyle("None", "")}
+class PromptStyle(typing.NamedTuple):
+    name: str
+    prompt: str
+    negative_prompt: str
 
-    if os.path.exists(filename):
-        with open(filename, "r", encoding="utf8", newline='') as file:
-            reader = csv.DictReader(file)
 
-            for row in reader:
-                res[row["name"]] = PromptStyle(row["name"], row["text"])
+def load_styles(path: str) -> dict[str, PromptStyle]:
+    styles = {"None": PromptStyle("None", "", "")}
 
-    return res
+    if os.path.exists(path):
+        with open(path, "r", encoding="utf8", newline='') as file:
+            reader = csv.DictReader(file)
+            for row in reader:
+                # Support loading old CSV format with "name, text"-columns
+                prompt = row["prompt"] if "prompt" in row else row["text"]
+                negative_prompt = row.get("negative_prompt", "")
+                styles[row["name"]] = PromptStyle(row["name"], prompt, negative_prompt)
 
+    return styles
 
-def apply_style_text(style_text, prompt):
-    if style_text == "":
-        return prompt
 
-    return prompt + ", " + style_text if prompt else style_text
+def merge_prompts(style_prompt: str, prompt: str) -> str:
+    parts = filter(None, (prompt.strip(), style_prompt.strip()))
+    return ", ".join(parts)
 
 
-def apply_style(p, style):
-    if type(p.prompt) == list:
-        p.prompt = [apply_style_text(style.text, x) for x in p.prompt]
+def apply_style(processing: StableDiffusionProcessing, style: PromptStyle) -> None:
+    if isinstance(processing.prompt, list):
+        processing.prompt = [merge_prompts(style.prompt, p) for p in processing.prompt]
     else:
-        p.prompt = apply_style_text(style.text, p.prompt)
+        processing.prompt = merge_prompts(style.prompt, processing.prompt)
 
-
-def save_style(filename, style):
-    with open(filename, "a", encoding="utf8", newline='') as file:
-        atstart = file.tell() == 0
-
-        writer = csv.DictWriter(file, fieldnames=["name", "text"])
-
-        if atstart:
-            writer.writeheader()
-
-        writer.writerow({"name": style.name, "text": style.text})
+    if isinstance(processing.negative_prompt, list):
+        processing.negative_prompt = [merge_prompts(style.negative_prompt, p) for p in processing.negative_prompt]
+    else:
+        processing.negative_prompt = merge_prompts(style.negative_prompt, processing.negative_prompt)
+
+
+def save_styles(path: str, styles: abc.Iterable[PromptStyle]) -> None:
+    # Write to temporary file first, so we don't nuke the file if something goes wrong
+    fd, temp_path = tempfile.mkstemp(".csv")
+    with os.fdopen(fd, "w", encoding="utf8", newline='') as file:
+        # _fields is actually part of the public API: typing.NamedTuple is a replacement for collections.NamedTuple,
+        # and collections.NamedTuple has explicit documentation for accessing _fields. Same goes for _asdict()
+        writer = csv.DictWriter(file, fieldnames=PromptStyle._fields)
+        writer.writeheader()
+        writer.writerows(style._asdict() for style in styles)
+
+    # Always keep a backup file around
+    shutil.copy(path, path + ".bak")
+    shutil.move(temp_path, path)
-- 
cgit v1.2.3


From cacd14bee89d1c1ff00fb11ebd2c68b407c90f1f Mon Sep 17 00:00:00 2001
From: cryzed <cryzed@googlemail.com>
Date: Sun, 11 Sep 2022 20:17:09 +0200
Subject: Only create backup if path exists

---
 modules/styles.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

(limited to 'modules/styles.py')

diff --git a/modules/styles.py b/modules/styles.py
index 25ec504b..bc7f070f 100644
--- a/modules/styles.py
+++ b/modules/styles.py
@@ -63,5 +63,6 @@ def save_styles(path: str, styles: abc.Iterable[PromptStyle]) -> None:
         writer.writerows(style._asdict() for style in styles)
 
     # Always keep a backup file around
-    shutil.copy(path, path + ".bak")
+    if os.path.exists(path):
+        shutil.move(path, path + ".bak")
     shutil.move(temp_path, path)
-- 
cgit v1.2.3