From 86867e153f4449167e3489323df35cf04f1fffa0 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Fri, 9 Sep 2022 23:16:02 +0300
Subject: support for prompt styles fix broken prompt matrix

---
 modules/processing.py | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

(limited to 'modules/processing.py')

diff --git a/modules/processing.py b/modules/processing.py
index 09680fbf..7e6cd8ee 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -16,6 +16,7 @@ from modules.shared import opts, cmd_opts, state
 import modules.shared as shared
 import modules.face_restoration
 import modules.images as images
+import modules.styles
 
 # some of those options should not be changed at all because they would break the model, so I removed them from options.
 opt_C = 4
@@ -29,13 +30,14 @@ def torch_gc():
 
 
 class StableDiffusionProcessing:
-    def __init__(self, sd_model=None, outpath_samples=None, outpath_grids=None, prompt="", seed=-1, subseed=-1, subseed_strength=0, seed_resize_from_h=-1, seed_resize_from_w=-1, sampler_index=0, batch_size=1, n_iter=1, steps=50, cfg_scale=7.0, width=512, height=512, restore_faces=False, tiling=False, do_not_save_samples=False, do_not_save_grid=False, extra_generation_params=None, overlay_images=None, negative_prompt=None):
+    def __init__(self, sd_model=None, outpath_samples=None, outpath_grids=None, prompt="", prompt_style=0, seed=-1, subseed=-1, subseed_strength=0, seed_resize_from_h=-1, seed_resize_from_w=-1, sampler_index=0, batch_size=1, n_iter=1, steps=50, cfg_scale=7.0, width=512, height=512, restore_faces=False, tiling=False, do_not_save_samples=False, do_not_save_grid=False, extra_generation_params=None, overlay_images=None, negative_prompt=None):
         self.sd_model = sd_model
         self.outpath_samples: str = outpath_samples
         self.outpath_grids: str = outpath_grids
         self.prompt: str = prompt
         self.prompt_for_display: str = None
         self.negative_prompt: str = (negative_prompt or "")
+        self.prompt_style: int = prompt_style
         self.seed: int = seed
         self.subseed: int = subseed
         self.subseed_strength: float = subseed_strength
@@ -154,8 +156,6 @@ def fix_seed(p):
 def process_images(p: StableDiffusionProcessing) -> Processed:
     """this is the main loop that both txt2img and img2img use; it calls func_init once inside all the scopes and func_sample once per batch"""
 
-    prompt = p.prompt
-
     assert p.prompt is not None
     torch_gc()
 
@@ -168,10 +168,12 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
 
     comments = []
 
-    if type(prompt) == list:
-        all_prompts = prompt
+    modules.styles.apply_style(p, shared.prompt_styles[p.prompt_style])
+
+    if type(p.prompt) == list:
+        all_prompts = p.prompt
     else:
-        all_prompts = p.batch_size * p.n_iter * [prompt]
+        all_prompts = p.batch_size * p.n_iter * [p.prompt]
 
     if type(p.seed) == list:
         all_seeds = int(p.seed)
@@ -207,7 +209,7 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
         
         negative_prompt_text = "\nNegative prompt: " + p.negative_prompt if p.negative_prompt else ""
 
-        return f"{p.prompt_for_display or prompt}{negative_prompt_text}\n{generation_params_text}".strip() + "".join(["\n\n" + x for x in comments])
+        return f"{all_prompts[index]}{negative_prompt_text}\n{generation_params_text}".strip() + "".join(["\n\n" + x for x in comments])
 
     if os.path.exists(cmd_opts.embeddings_dir):
         model_hijack.load_textual_inversion_embeddings(cmd_opts.embeddings_dir, p.sd_model)
-- 
cgit v1.2.3