This fixes a server crash (TypeError NoneType in torch/nn/modules/module.py) when the token counter API is triggered during model load/switch.
Standard sequential null checks failed because the crash occurred inside PyTorch's internal attribute lookup process. The whole function is wrapped in try-except as a robust defense against this race condition.
The token counter was crashing during model loading/unloading (race condition) because it was attempting to access attributes on an object that was momentarily None.
This fix introduces a sequential and robust null check (sd_models.model_data is None or sd_models.model_data.sd_model is None) at the start of the function to safely return 0 tokens during unstable states.
The previous fix (v1) was incomplete, as sd_model.cond_stage_model could still be None during model loading/unloading sequences (observed with XYZ Plot model swapping).
This v2 fix adds a specific check for cond_stage_model to ensure stability in dynamic environments.
This fixes the critical `TypeError: argument of type 'NoneType' is not iterable` that frequently occurs in the token counter when the main Stable Diffusion model is being unloaded or reloaded (e.g., during XYZ Plot or manual model switching).
The error happens because the `get_token_counter` function tries to access `sd_models.model_data.sd_model.cond_stage_model` when `sd_model` is momentarily `None`.
The fix adds a robust check for `sd_models.model_data.sd_model is None` at the start of the function, ensuring the token counter safely returns 0 instead of causing an ASGI application crash across all environments (A1111, Forge/reForge).
This fixes the TypeError that occurs during model switching in sd-webui-reforge and possibly other optimized environments.
The error happens because 'get_token_counter' tries to access 'sd_models.model_data.sd_model.cond_stage_model' before the model is fully loaded (when 'sd_model' is None).
The fix adds a check for the 'None' value to safely return 0 tokens instead of throwing an error.