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/ui.py | 6 ++++++
 1 file changed, 6 insertions(+)

(limited to 'modules/ui.py')

diff --git a/modules/ui.py b/modules/ui.py
index 5284a630..dcba8e88 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -152,6 +152,12 @@ def connect_clear_prompt(button):
 
 
 def update_token_counter(text, steps, styles, *, is_positive=True):
+    params = script_callbacks.BeforeTokenCounterParams(text, steps, styles, is_positive=is_positive)
+    script_callbacks.before_token_counter_callback(params)
+    text = params.prompt
+    steps = params.steps
+    styles = params.styles
+    is_positive = params.is_positive
 
     if shared.opts.include_styles_into_token_counters:
         apply_styles = shared.prompt_styles.apply_styles_to_prompt if is_positive else shared.prompt_styles.apply_negative_styles_to_prompt
-- 
cgit v1.2.3