Add support for saving styles with negative prompts

.gitignore

@@ -12,5 +12,6 @@ __pycache__
 /webui.settings.bat
 /embeddings
 /styles.csv
+/styles.csv.bak
 /webui-user.bat
 /interrogate

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)

modules/ui.py

@@ -228,17 +228,17 @@ def create_seed_inputs():
     return seed, subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w
 
 
-def add_style(style_name, text):
-    if style_name is None:
+def add_style(name: str, prompt: str, negative_prompt: str):
+    if name is None:
         return [gr_show(), gr_show()]
 
-    style = modules.styles.PromptStyle(style_name, text)
-
-    modules.styles.save_style(shared.styles_filename, style)
-
+    style = modules.styles.PromptStyle(name, prompt, negative_prompt)
     shared.prompt_styles[style.name] = style
+    # Save all loaded prompt styles: this allows us to update the storage format in the future more easily, because we
+    # reserialize all styles every time we save them
+    modules.styles.save_styles(shared.styles_filename, shared.prompt_styles.values())
 
-    update = {"visible": True, "choices": [k for k, v in shared.prompt_styles.items()], "__type__": "update"}
+    update = {"visible": True, "choices": list(shared.prompt_styles), "__type__": "update"}
     return [update, update]
 
 
@@ -251,7 +251,7 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
     with gr.Blocks(analytics_enabled=False) as txt2img_interface:
         with gr.Row(elem_id="toprow"):
             txt2img_prompt = gr.Textbox(label="Prompt", elem_id="txt2img_prompt", show_label=False, placeholder="Prompt", lines=1)
-            negative_prompt = gr.Textbox(label="Negative prompt", elem_id="txt2img_negative_prompt", show_label=False, placeholder="Negative prompt", lines=1)
+            txt2img_negative_prompt = gr.Textbox(label="Negative prompt", elem_id="txt2img_negative_prompt", show_label=False, placeholder="Negative prompt", lines=1)
             txt2img_prompt_style = gr.Dropdown(label="Style", show_label=False, elem_id="style_index", choices=[k for k, v in shared.prompt_styles.items()], value=next(iter(shared.prompt_styles.keys())), visible=len(shared.prompt_styles) > 1)
             roll = gr.Button('Roll', elem_id="txt2img_roll", visible=len(shared.artist_db.artists) > 0)
             submit = gr.Button('Generate', elem_id="txt2img_generate", variant='primary')
@@ -308,7 +308,7 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
                 _js="submit",
                 inputs=[
                     txt2img_prompt,
-                    negative_prompt,
+                    txt2img_negative_prompt,
                     txt2img_prompt_style,
                     steps,
                     sampler_index,
@@ -372,7 +372,7 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
     with gr.Blocks(analytics_enabled=False) as img2img_interface:
         with gr.Row(elem_id="toprow"):
             img2img_prompt = gr.Textbox(label="Prompt", elem_id="img2img_prompt", show_label=False, placeholder="Prompt", lines=1)
-            negative_prompt = gr.Textbox(label="Negative prompt", elem_id="img2img_negative_prompt", show_label=False, placeholder="Negative prompt", lines=1)
+            img2img_negative_prompt = gr.Textbox(label="Negative prompt", elem_id="img2img_negative_prompt", show_label=False, placeholder="Negative prompt", lines=1)
             img2img_prompt_style = gr.Dropdown(label="Style", show_label=False, elem_id="style_index", choices=[k for k, v in shared.prompt_styles.items()], value=next(iter(shared.prompt_styles.keys())), visible=len(shared.prompt_styles) > 1)
             img2img_interrogate = gr.Button('Interrogate', elem_id="img2img_interrogate", variant='primary')
             submit = gr.Button('Generate', elem_id="img2img_generate", variant='primary')
@@ -441,7 +441,6 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
                         img2img_save_style = gr.Button('Save prompt as style')
 
                 progressbar = gr.HTML(elem_id="progressbar")
-                style_dummpy = gr.Textbox(visible=False)
 
                 with gr.Group():
                     html_info = gr.HTML()
@@ -510,7 +509,7 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
                 _js="submit",
                 inputs=[
                     img2img_prompt,
-                    negative_prompt,
+                    img2img_negative_prompt,
                     img2img_prompt_style,
                     init_img,
                     init_img_with_mask,
@@ -580,11 +579,14 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
                 ]
             )
 
-            for button, propmt in zip([txt2img_save_style, img2img_save_style], [txt2img_prompt, img2img_prompt]):
+            dummy_component = gr.Label(visible=False)
+            for button, (prompt, negative_prompt) in zip([txt2img_save_style, img2img_save_style], [(txt2img_prompt, txt2img_negative_prompt), (img2img_prompt, img2img_negative_prompt)]):
                 button.click(
                     fn=add_style,
                     _js="ask_for_style_name",
-                    inputs=[style_dummpy, propmt],
+                    # Have to pass empty dummy component here, because the JavaScript and Python function have to accept
+                    # the same number of parameters, but we only know the style-name after the JavaScript prompt
+                    inputs=[dummy_component, prompt, negative_prompt],
                     outputs=[txt2img_prompt_style, img2img_prompt_style],
                 )
 

script.js

@@ -186,11 +186,7 @@ window.addEventListener('paste', e => {
         });
 });
 
-function ask_for_style_name(style_name, text){
-    input = prompt('Style name:');
-    if (input === null) {
-        return [null, null]
-    }
-
-    return [input, text]
+function ask_for_style_name(_, prompt_text, negative_prompt_text) {
+    name_ = prompt('Style name:')
+    return name_ === null ? [null, null, null]: [name_, prompt_text, negative_prompt_text]
 }