Update my_model/utilities/state_manager.py

my_model/utilities/state_manager.py

@@ -7,7 +7,7 @@ from my_model.KBVQA import KBVQA, prepare_kbvqa_model
 
 class StateManager:
     def __init__(self):
-        self.initialize_state()
+       # self.initialize_state()
 
     def initialize_state(self):
         if 'images_data' not in st.session_state:
@@ -46,15 +46,15 @@ class StateManager:
         df = pd.DataFrame(data)
         st.table(df)
 
-    def load_model(self, detection_model, confidence_level):
+    def load_model(self):
         """Load the KBVQA model with specified settings."""
         try:
             free_gpu_resources()
-            st.text("Loading the model, please wait...")
-            st.session_state['kbvqa'] = prepare_kbvqa_model(detection_model)
-            st.session_state['kbvqa'].detection_confidence = confidence_level
-            self.update_model_settings(detection_model, confidence_level)
-            st.write("Model is ready for inference.")
+            st.text("Loading the model, this should take no more than a few minutes, please wait...")
+            st.session_state['kbvqa'] = prepare_kbvqa_model(st.state_session.detection_model)
+            st.session_state['kbvqa'].detection_confidence = st.state_session.confidence_level
+            #self.update_model_settings(detection_model, confidence_level)
+            st.text("Model is ready for inference.")
             free_gpu_resources()
         except Exception as e:
             st.error(f"Error loading model: {e}")