Update my_model/state_manager.py

my_model/state_manager.py

@@ -41,22 +41,22 @@ class StateManager:
         """
         Sets up user interface widgets for selecting models, settings, and displaying model settings conditionally.
         """
-        disabled=st.session_state['loading_in_progress']
-        self.col1.selectbox("Choose a method:", ["Fine-Tuned Model", "In-Context Learning (n-shots)"], index=0, key='method', disabled=self.is_widget_disabled)
-        detection_model = self.col1.selectbox("Choose a model for objects detection:", ["yolov5", "detic"], index=1, key='detection_model', disabled=self.is_widget_disabled)
+        
+        self.col1.selectbox("Choose a method:", ["Fine-Tuned Model", "In-Context Learning (n-shots)"], index=0, key='method', on_click=self.disable_widgets, disabled=self.is_widget_disabled)
+        detection_model = self.col1.selectbox("Choose a model for objects detection:", ["yolov5", "detic"], index=1, key='detection_model', on_click=self.disable_widgets, disabled=self.is_widget_disabled)
         default_confidence = 0.2 if st.session_state.detection_model == "yolov5" else 0.4
-        self.set_slider_value(text="Select minimum detection confidence level", min_value=0.1, max_value=0.9, value=default_confidence, step=0.1, slider_key_name='confidence_level', col=self.col1, disabled=self.is_widget_disabled)
+        self.set_slider_value(text="Select minimum detection confidence level", min_value=0.1, max_value=0.9, value=default_confidence, step=0.1, slider_key_name='confidence_level', col=self.col1)
 
         # Conditional display of model settings
 
  
-        show_model_settings = self.col3.checkbox("Show Model Settings", False)
+        show_model_settings = self.col3.checkbox("Show Model Settings", False, on_click=self.disable_widgets, disabled=self.is_widget_disabled)
         if show_model_settings:
             self.display_model_settings()
 
 
         
-    def set_slider_value(self, text, min_value, max_value, value, step, slider_key_name, col=None, disabled=False):
+    def set_slider_value(self, text, min_value, max_value, value, step, slider_key_name, col=None):
         """
         Creates a slider widget with the specified parameters, optionally placing it in a specific column.
 
@@ -71,9 +71,9 @@ class StateManager:
         """
         
         if col is None:
-            return st.slider(text, min_value, max_value, value, step, key=slider_key_name, disabled=disabled)
+            return st.slider(text, min_value, max_value, value, step, key=slider_key_name, on_click=self.disable_widgets, disabled=self.is_widget_disabledd)
         else:
-            return col.slider(text, min_value, max_value, value, step, key=slider_key_name, disabled=disabled)
+            return col.slider(text, min_value, max_value, value, step, key=slider_key_name, on_click=self.disable_widgets, disabled=self.is_widget_disabled)
 
     @property
     def is_widget_disabled(self):
@@ -129,14 +129,13 @@ class StateManager:
         """
         
         try:
-            st.session_state['loading_in_progress'] = True
+
             free_gpu_resources()
             st.session_state['kbvqa'] = prepare_kbvqa_model()
             st.session_state['kbvqa'].detection_confidence = st.session_state.confidence_level
             # 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
             st.session_state['button_label'] = "Reload Model"
             free_gpu_resources()
             
@@ -145,7 +144,6 @@ class StateManager:
 
     def force_reload_model(self):
         try:
-            st.session_state['loading_in_progress'] = True
             self.delete_model()
             free_gpu_resources()
             st.session_state['kbvqa'] = prepare_kbvqa_model(force_reload=True)
@@ -153,7 +151,6 @@ 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}")
@@ -223,11 +220,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}")