Update my_model/state_manager.py

my_model/state_manager.py

@@ -139,6 +139,7 @@ class StateManager:
 
     def force_reload_model(self):
         try:
+            st.session_state['loading_in_progress'] = True
             self.delete_model()
             self.initialize_state()
             free_gpu_resources()
@@ -147,6 +148,7 @@ class StateManager:
             # Update the previous state with current session state values
             st.session_state['previous_state'] = {'method': st.session_state.method, 'detection_model': st.session_state.detection_model, 'confidence_level': st.session_state.confidence_level}
             st.session_state['model_loaded'] = True
+            st.session_state['loading_in_progress'] = False
             free_gpu_resources()
         except Exception as e:
             st.error(f"Error reloading model: {e}")
@@ -216,9 +218,11 @@ class StateManager:
         try:
             free_gpu_resources()
             if self.is_model_loaded():
+                st.session_state['loading_in_progress'] = True
                 prepare_kbvqa_model(only_reload_detection_model=True)
                 st.session_state['kbvqa'].detection_confidence = st.session_state.confidence_level
                 self.col1.success("Model reloaded with updated settings and ready for inference.")
+                st.session_state['loading_in_progress'] = False
             free_gpu_resources()
         except Exception as e:
             st.error(f"Error reloading detection model: {e}")