diff --git a/scripts/physton_prompt/get_token_counter.py b/scripts/physton_prompt/get_token_counter.py
index 2d005c1..462fc5c 100644
--- a/scripts/physton_prompt/get_token_counter.py
+++ b/scripts/physton_prompt/get_token_counter.py
@@ -4,10 +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:
+    # FIX: Robust Null Check to prevent TypeError during model loading/unloading.
+    # Checks for the existence of model_data and its property sd_model sequentially.
+    if sd_models.model_data is None or sd_models.model_data.sd_model is None:
         return {"token_count": 0, "max_length": 0}
 
     # copy from modules.ui.py