Update my_model/state_manager.py

my_model/state_manager.py

@@ -22,8 +22,6 @@ class StateManager:
             st.session_state['kbvqa'] = None
         if "button_label" not in st.session_state: 
             st.session_state['button_label'] = "Load Model"
-        if "previous_state" not in st.session_state:
-            st.session_state['previous_state'] = {}
         if 'loading_in_progress' not in st.session_state:
             st.session_state['loading_in_progress'] = False
         if 'load_button_clicked' not in st.session_state:
@@ -38,6 +36,8 @@ class StateManager:
             st.session_state['settings_changed'] = self.settings_changed
         if 'model_loaded' not in st.session_state:
             st.session_state['model_loaded'] = self.is_model_loaded
+        if "previous_state" not in st.session_state:
+            st.session_state['previous_state'] = {'method': None, 'detection_model': None, 'confidence_level': None}
 
     
 
@@ -120,7 +120,10 @@ class StateManager:
         df = pd.DataFrame(data).reset_index(drop=True)
         st.write(df)
         
-
+    def update_prev_state(self):
+        for key in st.session_state['previous_state']:
+            st.session_state['previous_state'][key] = st.session_state[key]
+            
     def load_model(self):
         """
         Loads the KBVQA model based on the chosen method and settings.
@@ -137,7 +140,7 @@ class StateManager:
             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}
+            self.update_prev_state()
             st.session_state['model_loaded'] = True
             st.session_state['button_label'] = "Reload Model"
             free_gpu_resources()
@@ -152,7 +155,8 @@ class StateManager:
             st.session_state['kbvqa'] = prepare_kbvqa_model(force_reload=True)
             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}
+            self.update_prev_state()
+            
             st.session_state['model_loaded'] = True
             free_gpu_resources()
         except Exception as e:
@@ -184,6 +188,8 @@ class StateManager:
             bool: True if any change is found, False otherwise.
         """
         for key in st.session_state['previous_state']:
+            if key == 'confidence_level':
+                continue  # confidence_level tracker is separate
             if st.session_state[key] != st.session_state['previous_state'][key]:
                 return True  # Found a change
         else: return False  # No changes found   
@@ -222,14 +228,10 @@ class StateManager:
         try:
             free_gpu_resources()
             if self.is_model_loaded:
-          
-                st.write(st.session_state['detection_model'] == st.session_state['previous_state']['detection_model'])
-                st.write(st.session_state['method'] == st.session_state['previous_state']['method'])
-                    
                 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['previous_state'] = {'method': st.session_state.method, 'detection_model': st.session_state.detection_model}
+                self.update_prev_state
                 st.session_state['button_label'] = "Reload Model"
 
             free_gpu_resources()