fix(prompt): handle multi-modal processors in token counter

For FLUX2 models using Mistral3/PixtralProcessor, detect and use the
underlying text tokenizer instead of calling the processor directly,
which incorrectly tries to interpret text prompts as images.
This commit is contained in:
CalamitousFelicitousness
2025-12-10 18:54:39 +00:00
parent 886f57708c
commit a19c795b49
2 changed files with 32 additions and 15 deletions
+19 -10
View File
@@ -395,30 +395,39 @@ def get_prompt_schedule(prompt, steps):
def get_tokens(pipe, msg, prompt):
global token_dict, token_type # pylint: disable=global-statement
token_count = 0
if shared.sd_loaded and hasattr(pipe, 'tokenizer') and pipe.tokenizer is not None:
tokenizer = pipe.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
prompt = prompt.replace(' BOS ', ' !!!!!!!! ').replace(' EOS ', ' !!!!!!! ')
debug(f'Prompt tokenizer: type={msg} prompt="{prompt}"')
if token_dict is None or token_type != shared.sd_model_type:
token_type = shared.sd_model_type
fn = pipe.tokenizer.name_or_path
fn = getattr(tokenizer, 'name_or_path', '')
if fn.endswith('tokenizer'):
fn = os.path.join(pipe.tokenizer.name_or_path, 'vocab.json')
fn = os.path.join(fn, 'vocab.json')
else:
fn = os.path.join(pipe.tokenizer.name_or_path, 'tokenizer', 'vocab.json')
fn = os.path.join(fn, 'tokenizer', 'vocab.json')
token_dict = shared.readfile(fn, silent=True)
for k, v in pipe.tokenizer.added_tokens_decoder.items():
added_tokens = getattr(tokenizer, 'added_tokens_decoder', {})
for k, v in added_tokens.items():
token_dict[str(v)] = k
shared.log.debug(f'Tokenizer: words={len(token_dict)} file="{fn}"')
has_bos_token = pipe.tokenizer.bos_token_id is not None
has_eos_token = pipe.tokenizer.eos_token_id is not None
ids = pipe.tokenizer(prompt)
ids = getattr(ids, 'input_ids', [])
has_bos_token = getattr(tokenizer, 'bos_token_id', None) is not None
has_eos_token = getattr(tokenizer, 'eos_token_id', None) is not None
try:
ids = tokenizer(prompt)
ids = getattr(ids, 'input_ids', [])
except Exception:
ids = []
if has_bos_token and has_eos_token:
for i in range(len(ids)):
if ids[i] == 21622:
ids[i] = pipe.tokenizer.bos_token_id
ids[i] = tokenizer.bos_token_id
elif ids[i] == 15203:
ids[i] = pipe.tokenizer.eos_token_id
ids[i] = tokenizer.eos_token_id
tokens = []
for i in ids:
try: