Update my_model/state_manager.py

my_model/state_manager.py

@@ -29,6 +29,10 @@ class StateManager:
             st.session_state['model_loaded'] = False
         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:
+            st.session_state['load_button_clicked'] = False
+        if 'force_reload_button_clicked' not in st.session_state:
+            st.session_state['force_reload_button_clicked'] = False
 
             
 
@@ -132,6 +136,8 @@ class StateManager:
 
     def force_reload_model(self):
         try:
+            self.delete_model()
+            self.initialize_state()
             free_gpu_resources()
             st.session_state['kbvqa'] = prepare_kbvqa_model(force_reload=True)
             st.session_state['kbvqa'].detection_confidence = st.session_state.confidence_level
@@ -140,9 +146,22 @@ class StateManager:
             st.session_state['model_loaded'] = True
             free_gpu_resources()
         except Exception as e:
-            st.error(f"Error loading model: {e}")
+            st.error(f"Error reloading model: {e}")
             free_gpu_resources()
         
+    def delete_model(self):
+        """
+        Forces a reload of all models, freeing up GPU resources. This method deletes the current models and calls `free_gpu_resources`.
+        """
+        
+        free_gpu_resources()
+        
+        if st.session_state['kbvqa'] is not None:
+            del st.session_state['kbvqa']
+            st.write('Model deleted')
+            free_gpu_resources()
+        free_gpu_resources()
+
         
     # Function to check if any session state values have changed
     def has_state_changed(self):