diff --git a/modules/ui_common.py b/modules/ui_common.py
index 3b7530e08..133136d8d 100644
--- a/modules/ui_common.py
+++ b/modules/ui_common.py
@@ -1,9 +1,11 @@
import json
import html
import os
+import re
import shutil
import platform
import subprocess
+from weakref import WeakSet
import gradio as gr
from modules import paths, call_queue, shared, errors, ui_sections, ui_symbols, ui_components, generation_parameters_copypaste, images, scripts_manager, script_callbacks, infotext, processing
from modules.logger import log
@@ -13,6 +15,8 @@ folder_symbol = ui_symbols.folder
debug = log.trace if os.environ.get('SD_PASTE_DEBUG', None) is not None else lambda *args, **kwargs: None
debug('Trace: PASTE')
+warn_once_set = WeakSet()
+
def gr_show(visible=True):
return {"visible": visible, "__type__": "update"}
@@ -427,32 +431,52 @@ def connect_reuse_seed(seed: gr.Number, reuse_seed_btn: gr.Button, generation_in
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:
log.debug('Tokenizer busy')
- return f"{token_count}/{max_length}"
- 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="--/--", 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_attention == 'native':
+ 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)
- model_max_length = getattr(tokenizer, 'model_max_length', 0)
+ 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', 77)
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"{token_count}/{max_length}", visible=token_count > 0)
+
+ try:
+ try:
+ ids: list = getattr(tokenizer(prompt_list), 'input_ids', [])
+ except TypeError:
+ for p in prompt_list:
+ ids.append(getattr(tokenizer(p), 'input_ids', []))
+ except Exception as e:
+ if tokenizer not in warn_once_set:
+ log.warning(f"Token counter: {e}")
+ warn_once_set.add(tokenizer)
+ return gr.update(value=f"??/{max_length}", 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"{count_formatted}/{max_length}", visible=visible)