Merge pull request #369 from neon-aiart/fix/token-counter-v4

Fix: Use try-except to stabilize token counter during model access
pull/373/head
Physton 2025-11-26 20:12:09 +08:00 committed by GitHub
commit 596a931725
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 37 additions and 36 deletions

View File

@ -4,12 +4,9 @@ from functools import partial, reduce
def get_token_counter(text, steps):
# Check if the model is fully loaded to prevent TypeError during model switching.
# Checks both sd_model and its subcomponent (cond_stage_model).
if sd_models.model_data.sd_model is None or \
sd_models.model_data.sd_model.cond_stage_model is None:
return {"token_count": 0, "max_length": 0}
# FIX: Use try-except to safely handle PyTorch/model access errors (TypeError NoneType)
# that occur during model loading/switching when the token counter API is triggered.
try:
# copy from modules.ui.py
try:
text, _ = extra_networks.parse_prompt(text)
@ -41,3 +38,7 @@ def get_token_counter(text, steps):
key=lambda args: args[0])
return {"token_count": token_count, "max_length": max_length}
except Exception as e:
# return 0 token count if any error (model instability, parsing error, etc.) occurs during calculation
return {"token_count": 0, "max_length": 0}