Add optional detailed token counts

html/locale_en.json

@@ -1248,6 +1248,7 @@
       {"id":"txt2img_skip","label":"Skip","localized":"","hint":"Stop processing current job and continue processing","ui":"txt2img"},
       {"id":"txt2img_interrupt","label":"Stop","localized":"","hint":"Stop processing","ui":"txt2img"},
       {"id":"","label":"Show","localized":"","hint":"Show image location","ui":"txt2img"},
+      {"id":"","label":"Show detailed token counts","localized":"","hint":"Show token counts split by \"BREAK\" segment markers (and by line breaks if that setting is enabled.","ui":"settings_text_encoder"},
       {"id":"","label":"Save","localized":"","hint":"Save image","ui":"txt2img"},
       {"id":"","label":"Sketch","localized":"","hint":"","ui":"img2img"},
       {"id":"video_params_size","label":"Size & Inputs","localized":"","hint":"Settings related to generation resolution and additional input media","ui":"video"},

modules/shared.py

@@ -310,6 +310,7 @@ def list_samplers():
     "sd_textencoder_cache_size": OptionInfo(4, "Text encoder cache size", gr.Slider, {"minimum": 0, "maximum": 16, "step": 1}),
     "sd_textencder_linebreak": OptionInfo(True, "Use line break as prompt segment marker", gr.Checkbox),
     "diffusers_zeros_prompt_pad": OptionInfo(False, "Use zeros for prompt padding", gr.Checkbox),
+    "prompt_detailed_tokens": OptionInfo(False, "Show detailed token counts", gr.Checkbox),
     "te_optional_sep": OptionInfo("<h2>Optional</h2>", "", gr.HTML),
     "te_shared_t5": OptionInfo(True, "T5: Use shared instance of text encoder"),
     "te_pooled_embeds": OptionInfo(False, "SDXL: Use weighted pooled embeds"),

modules/ui_common.py

@@ -1,6 +1,7 @@
 import json
 import html
 import os
+import re
 import shutil
 import platform
 import subprocess
@@ -422,32 +423,50 @@ def copy_seed(gen_info_string: str, index: int):
         reuse_seed_btn.click(fn=copy_seed, _js="(x, y) => [x, selected_gallery_index()]", show_progress='hidden', inputs=[generation_info, dummy_component], outputs=[seed, dummy_component, subseed_strength])
 
 
-def update_token_counter(text):
-    token_count = 0
-    max_length = 75
+def update_token_counter(text: str):
     if shared.state.job_count > 0:
         shared.log.debug('Tokenizer busy')
-        return f"<span class='gr-box gr-text-input'>{token_count}/{max_length}</span>"
-    from modules import extra_networks
-    if isinstance(text, list):
-        prompt, _ = extra_networks.parse_prompts(text)
-    else:
-        prompt, _ = extra_networks.parse_prompt(text)
+        return gr.update(value="<span class='gr-box gr-text-input'>--/--</span>", visible=True)
+
+    from modules.extra_networks import parse_prompt
+
+    count_formatted = '0'
+    visible = False
+
+    prompt, _ = parse_prompt(text)
+    prompt_list = [prompt]
+    ids = []
     if shared.sd_loaded and hasattr(shared.sd_model, 'tokenizer') and shared.sd_model.tokenizer is not None:
+        if shared.opts.prompt_detailed_tokens:
+            p_split = re.compile(r'\bBREAK\b|\n' if shared.opts.sd_textencder_linebreak else r'\bBREAK\b')
+            prompt_list = re.split(p_split, prompt)
+
         tokenizer = shared.sd_model.tokenizer
         # For multi-modal processors (e.g., PixtralProcessor), use the underlying text tokenizer
         if hasattr(tokenizer, 'tokenizer') and tokenizer.tokenizer is not None:
             tokenizer = tokenizer.tokenizer
-        has_bos_token = hasattr(tokenizer, 'bos_token_id') and tokenizer.bos_token_id is not None
-        has_eos_token = hasattr(tokenizer, 'eos_token_id') and tokenizer.eos_token_id is not None
-        try:
-            ids = tokenizer(prompt)
-            ids = getattr(ids, 'input_ids', [])
-        except Exception:
-            ids = []
-        token_count = len(ids) - int(has_bos_token) - int(has_eos_token)
+        has_bos_token = getattr(tokenizer, 'bos_token_id', None) is not None
+        has_eos_token = getattr(tokenizer, 'eos_token_id', None) is not None
         model_max_length = getattr(tokenizer, 'model_max_length', 0)
         max_length = model_max_length - int(has_bos_token) - int(has_eos_token)
         if max_length is None or max_length < 0 or max_length > 10000:
             max_length = 0
-    return gr.update(value=f"<span class='gr-box gr-text-input'>{token_count}/{max_length}</span>", visible=token_count > 0)
+
+        try:
+            try:
+                ids = getattr(tokenizer(prompt_list), 'input_ids', [])
+            except TypeError:
+                for p in prompt_list:
+                    ids.append(getattr(tokenizer(p), 'input_ids', []))
+        except Exception as e:
+            shared.log.warning("Token counter:", e)
+            return gr.update(value=f"<span class='gr-box gr-text-input'>??/{max_length}</span>", visible=True)
+
+        token_counts = [len(group) - int(has_bos_token) - int(has_eos_token) for group in ids]
+        if len(token_counts) > 1:
+            visible = True
+            count_formatted = f"{token_counts} {sum(token_counts)}" if shared.opts.prompt_detailed_tokens else str(sum(token_counts))
+        elif len(token_counts) == 1 and token_counts[0] > 0:
+            visible = True
+            count_formatted = str(token_counts[0])
+    return gr.update(value=f"<span class='gr-box gr-text-input'>{count_formatted}/{max_length}</span>", visible=visible)