From b7f45e67dcf48914d2f34d4ace977a431a5aa12e Mon Sep 17 00:00:00 2001
From: AUTOMATIC1111 <16777216c@gmail.com>
Date: Sun, 11 Feb 2024 12:56:53 +0300
Subject: add before_token_counter callback and use it for prompt comments

---
 modules/processing_scripts/comments.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

(limited to 'modules/processing_scripts')

diff --git a/modules/processing_scripts/comments.py b/modules/processing_scripts/comments.py
index 316356c7..638e39f2 100644
--- a/modules/processing_scripts/comments.py
+++ b/modules/processing_scripts/comments.py
@@ -1,4 +1,4 @@
-from modules import scripts, shared
+from modules import scripts, shared, script_callbacks
 import re
 
 
@@ -27,6 +27,16 @@ class ScriptStripComments(scripts.Script):
         p.main_negative_prompt = strip_comments(p.main_negative_prompt)
 
 
+def before_token_counter(params: script_callbacks.BeforeTokenCounterParams):
+    if not shared.opts.enable_prompt_comments:
+        return
+
+    params.prompt = strip_comments(params.prompt)
+
+
+script_callbacks.on_before_token_counter(before_token_counter)
+
+
 shared.options_templates.update(shared.options_section(('sd', "Stable Diffusion", "sd"), {
     "enable_prompt_comments": shared.OptionInfo(True, "Enable comments").info("Use # anywhere in the prompt to hide the text between # and the end of the line from the generation."),
 }))
-- 
cgit v1.2.3