From a20103a86c6d50056c5291a55ba5e05c5aad3632 Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Mon, 10 Jun 2024 15:48:55 -0400
Subject: [PATCH] add prompt attention normalization

---
 CHANGELOG.md                       |  5 +++++
 extensions-builtin/sdnext-modernui |  2 +-
 modules/prompt_parser_diffusers.py | 22 ++++++++++++++++++++--
 modules/shared.py                  |  2 +-
 webui.py                           |  3 +--
 5 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 82704c6e0..f03741b59 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -52,6 +52,9 @@
 
 ### Improvements
 
+- additional modernui themes
+- reintroduce prompt attention normalization, disabled by default, enable in settings -> execution  
+  this can drastically help with unbalanced prompts  
 - further work on improving python 3.12 functionality and remove experimental flag  
   note: recommended version remains python 3.11 for all users except if you're using directml and then its python 3.10  
 - improved **installer** for initial installs  
@@ -65,6 +68,7 @@
   add `cli/image-exif.py` that can be used to view/extract metadata from images  
 - lower overhead on generate calls  
 - auto-synchronize modernui and core branches  
+- add option to pad prompt with empty embeds, thanks @Disty
 
 ## Fixes
 
@@ -72,6 +76,7 @@
 - fix apply/unapply hidiffusion for sd15  
 - fix controlnet reference enabled check  
 - fix face-hires with control batch count  
+- apply rollback-vae option to latest torch versions, thanks @Iaotle
 
 ## Update for 2024-06-02
 
diff --git a/extensions-builtin/sdnext-modernui b/extensions-builtin/sdnext-modernui
index cc2e7ee98..285743a83 160000
--- a/extensions-builtin/sdnext-modernui
+++ b/extensions-builtin/sdnext-modernui
@@ -1 +1 @@
-Subproject commit cc2e7ee980be3efaa514c68fac2b715cddfbc072
+Subproject commit 285743a83f251ae23e3a4120d15badcead4eab33
diff --git a/modules/prompt_parser_diffusers.py b/modules/prompt_parser_diffusers.py
index 426003bfb..8df78d3bc 100644
--- a/modules/prompt_parser_diffusers.py
+++ b/modules/prompt_parser_diffusers.py
@@ -183,12 +183,30 @@ def encode_prompts(pipe, p, prompts: list, negative_prompts: list, steps: int, c
         return
 
 
+def normalize_prompt(pairs: list):
+    num_words = 0
+    total_weight = 0
+    for section in pairs:
+        words = len(section[0].split())
+        if section[1] == -1: # control tokens
+            continue
+        num_words += words
+        total_weight += section[1] * words
+    avg_weight = round(100 * total_weight / num_words) / 100 if num_words > 0 else 1
+    debug(f'Prompt stats: words={num_words} weight={avg_weight}')
+    for section in pairs:
+        section[1] = section[1] / avg_weight if section[1] != -1 else -1 # skip control tokens
+    debug(f'Prompt normalized: {pairs}')
+    return pairs
+
+
 def get_prompts_with_weights(prompt: str):
     t0 = time.time()
-    manager = DiffusersTextualInversionManager(shared.sd_model,
-                                               shared.sd_model.tokenizer or shared.sd_model.tokenizer_2)
+    manager = DiffusersTextualInversionManager(shared.sd_model, shared.sd_model.tokenizer or shared.sd_model.tokenizer_2)
     prompt = manager.maybe_convert_prompt(prompt, shared.sd_model.tokenizer or shared.sd_model.tokenizer_2)
     texts_and_weights = prompt_parser.parse_prompt_attention(prompt)
+    if shared.opts.prompt_mean_norm:
+        texts_and_weights = normalize_prompt(texts_and_weights)
     texts, text_weights = zip(*texts_and_weights)
     debug(f'Prompt: weights={texts_and_weights} time={(time.time() - t0):.3f}')
     return texts, text_weights
diff --git a/modules/shared.py b/modules/shared.py
index ce5f82b4e..5c430c309 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -395,7 +395,7 @@ options_templates.update(options_section(('sd', "Execution & Models"), {
     "stream_load": OptionInfo(False, "Load models using stream loading method", gr.Checkbox, {"visible": not native }),
     "model_reuse_dict": OptionInfo(False, "Reuse loaded model dictionary", gr.Checkbox, {"visible": False}),
     "prompt_attention": OptionInfo("Full parser", "Prompt attention parser", gr.Radio, {"choices": ["Full parser", "Compel parser", "A1111 parser", "Fixed attention"] }),
-    "prompt_mean_norm": OptionInfo(True, "Prompt attention normalization", gr.Checkbox, {"visible": not native }),
+    "prompt_mean_norm": OptionInfo(False, "Prompt attention normalization", gr.Checkbox),
     "comma_padding_backtrack": OptionInfo(20, "Prompt padding", gr.Slider, {"minimum": 0, "maximum": 74, "step": 1, "visible": not native }),
     "sd_checkpoint_cache": OptionInfo(0, "Cached models", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1, "visible": not native }),
     "sd_vae_checkpoint_cache": OptionInfo(0, "Cached VAEs", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1, "visible": False}),
diff --git a/webui.py b/webui.py
index c0fdd7235..43eacb944 100644
--- a/webui.py
+++ b/webui.py
@@ -60,7 +60,6 @@ fastapi_args = {
 }
 
 import modules.sd_hijack
-from packaging.version import Version
 timer.startup.record("ldm")
 modules.loader.initialized = True
 
@@ -69,7 +68,7 @@ def check_rollback_vae():
         if not torch.cuda.is_available():
             log.error("Rollback VAE functionality requires compatible GPU")
             shared.cmd_opts.rollback_vae = False
-        elif not Version(torch.__version__) >= Version("2.1"):
+        elif torch.__version__.startswith('1.') or torch.__version__.startswith('2.0'):
             log.error("Rollback VAE functionality requires Torch 2.1 or higher")
             shared.cmd_opts.rollback_vae = False
         elif 0 < torch.cuda.get_device_capability()[0] < 8: