Count tokens of enabled styles

1 files changed, 16 insertions, 7 deletions

diff --git a/modules/ui.py b/modules/ui.py
index 9e90133b..5284a630 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -151,7 +151,12 @@ def connect_clear_prompt(button):
     )

 

 

-def update_token_counter(text, steps, *, is_positive=True):

+def update_token_counter(text, steps, styles, *, is_positive=True):

+

+    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

+        text = apply_styles(text, styles)

+

     try:

         text, _ = extra_networks.parse_prompt(text)

 

@@ -173,8 +178,8 @@ def update_token_counter(text, steps, *, is_positive=True):
     return f"<span class='gr-box gr-text-input'>{token_count}/{max_length}</span>"

 

 

-def update_negative_prompt_token_counter(text, steps):

-    return update_token_counter(text, steps, is_positive=False)

+def update_negative_prompt_token_counter(*args):

+    return update_token_counter(*args, is_positive=False)

 

 

 def setup_progressbar(*args, **kwargs):

@@ -485,8 +490,10 @@ def create_ui():
                 height,

             ]

 

-            toprow.token_button.click(fn=wrap_queued_call(update_token_counter), inputs=[toprow.prompt, steps], outputs=[toprow.token_counter])

-            toprow.negative_token_button.click(fn=wrap_queued_call(update_negative_prompt_token_counter), inputs=[toprow.negative_prompt, steps], outputs=[toprow.negative_token_counter])

+            toprow.ui_styles.dropdown.change(fn=wrap_queued_call(update_token_counter), inputs=[toprow.prompt, steps, toprow.ui_styles.dropdown], outputs=[toprow.token_counter])

+            toprow.ui_styles.dropdown.change(fn=wrap_queued_call(update_negative_prompt_token_counter), inputs=[toprow.negative_prompt, steps, toprow.ui_styles.dropdown], outputs=[toprow.negative_token_counter])

+            toprow.token_button.click(fn=wrap_queued_call(update_token_counter), inputs=[toprow.prompt, steps, toprow.ui_styles.dropdown], outputs=[toprow.token_counter])

+            toprow.negative_token_button.click(fn=wrap_queued_call(update_negative_prompt_token_counter), inputs=[toprow.negative_prompt, steps, toprow.ui_styles.dropdown], outputs=[toprow.negative_token_counter])

 

         extra_networks_ui = ui_extra_networks.create_ui(txt2img_interface, [txt2img_generation_tab], 'txt2img')

         ui_extra_networks.setup_ui(extra_networks_ui, output_panel.gallery)

@@ -824,8 +831,10 @@ def create_ui():
                 **interrogate_args,

             )

 

-            toprow.token_button.click(fn=update_token_counter, inputs=[toprow.prompt, steps], outputs=[toprow.token_counter])

-            toprow.negative_token_button.click(fn=wrap_queued_call(update_negative_prompt_token_counter), inputs=[toprow.negative_prompt, steps], outputs=[toprow.negative_token_counter])

+            toprow.ui_styles.dropdown.change(fn=wrap_queued_call(update_token_counter), inputs=[toprow.prompt, steps, toprow.ui_styles.dropdown], outputs=[toprow.token_counter])

+            toprow.ui_styles.dropdown.change(fn=wrap_queued_call(update_negative_prompt_token_counter), inputs=[toprow.negative_prompt, steps, toprow.ui_styles.dropdown], outputs=[toprow.negative_token_counter])

+            toprow.token_button.click(fn=update_token_counter, inputs=[toprow.prompt, steps, toprow.ui_styles.dropdown], outputs=[toprow.token_counter])

+            toprow.negative_token_button.click(fn=wrap_queued_call(update_negative_prompt_token_counter), inputs=[toprow.negative_prompt, steps, toprow.ui_styles.dropdown], outputs=[toprow.negative_token_counter])

 

             img2img_paste_fields = [

                 (toprow.prompt, "Prompt"),