Fixed the bug of model reloading

app.py

@@ -118,7 +118,6 @@ def _load_model_and_tokenizer_core(model_path_to_load, model_display_name_for_st
 
 def handle_model_selection_change(selected_model_name_ui):
     global MODEL, TOKENIZER, MASK_ID, CURRENT_MODEL_PATH, DEVICE, uni_prompting
-
     status_msg = ""
     # 初始化 Examples 的可见性更新
     vis_lm_base = gr.update(visible=False)
@@ -127,7 +126,6 @@ def handle_model_selection_change(selected_model_name_ui):
     vis_mmu_base = gr.update(visible=False)
     vis_mmu_mixcot = gr.update(visible=False)
     vis_mmu_max = gr.update(visible=False)
-
     # 根据选择的模型决定 thinking mode 的默认状态
     is_mixcot_model_selected = (selected_model_name_ui == "MMaDA-8B-MixCoT")
     
@@ -138,10 +136,8 @@ def handle_model_selection_change(selected_model_name_ui):
     
     lm_think_button_label = "Thinking Mode ✅" if current_thinking_mode_lm_state else "Thinking Mode ❌"
     mmu_think_button_label = "Thinking Mode ✅" if current_thinking_mode_mmu_state else "Thinking Mode ❌"
-
     update_think_button_lm = gr.update(value=lm_think_button_label)
     update_think_button_mmu = gr.update(value=mmu_think_button_label)
-
     if selected_model_name_ui == "MMaDA-8B-Max (coming soon)":
         MODEL = None
         TOKENIZER = None
@@ -169,17 +165,26 @@ def handle_model_selection_change(selected_model_name_ui):
             # 尝试加载模型
             status_msg = _load_model_and_tokenizer_core(actual_path, selected_model_name_ui)
             
-            # 检查模型是否成功加载
-            if "Error loading model" in status_msg or MODEL is None:
+            # 改进的错误检查逻辑
+            model_load_failed = False
+            
+            # 检查是否是明确的错误消息
+            if "Error loading model" in status_msg:
+                model_load_failed = True
+            # 检查模型是否为None，但排除"已经加载"的情况
+            elif MODEL is None:
+                if "is already loaded" not in status_msg:
+                    model_load_failed = True
+                    status_msg = f"Failed to properly load model '{selected_model_name_ui}'. {status_msg}"
+            
+            if model_load_failed:
                 # 如果是 MixCoT 模型但加载失败，则关闭 thinking mode
                 if is_mixcot_model_selected:
                     current_thinking_mode_lm_state = False
                     current_thinking_mode_mmu_state = False
                     update_think_button_lm = gr.update(value="Thinking Mode ❌")
                     update_think_button_mmu = gr.update(value="Thinking Mode ❌")
-                if MODEL is None and "Error" not in status_msg: # 补充一个通用错误信息
-                    status_msg = f"Failed to properly load model '{selected_model_name_ui}'. {status_msg}"
-            else: # 模型成功加载
+            else: # 模型成功加载或已经加载
                 if selected_model_name_ui == "MMaDA-8B-Base":
                     vis_lm_base = gr.update(visible=True)
                     vis_mmu_base = gr.update(visible=True)
@@ -187,7 +192,6 @@ def handle_model_selection_change(selected_model_name_ui):
                     vis_lm_mixcot = gr.update(visible=True)
                     vis_mmu_mixcot = gr.update(visible=True)
                     # thinking mode 已经在函数开头根据 is_mixcot_model_selected 设置为 True
-
     return (
         status_msg,
         vis_lm_base,