Update my_model/state_manager.py

my_model/state_manager.py

@@ -3,7 +3,7 @@ import copy
 from PIL import Image
 import streamlit as st
 from my_model.utilities.gen_utilities import free_gpu_resources
-from my_model.KBVQA import KBVQA, prepare_kbvqa_model, force_reload_model
+from my_model.KBVQA import KBVQA, prepare_kbvqa_model
 
 
 
@@ -130,16 +130,18 @@ class StateManager:
         except Exception as e:
             st.error(f"Error loading model: {e}")
 
-    def force_reload_kbvqa(self):
+    def force_reload_model(self):
         try:
             free_gpu_resources()
-            st.session_state['kbvqa'] = force_reload_model()
+            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}
             st.session_state['model_loaded'] = True
+            free_gpu_resources()
         except Exception as e:
             st.error(f"Error loading model: {e}")
+            free_gpu_resources()
         
         
     # Function to check if any session state values have changed