Update my_model/utilities/state_manager.py

my_model/utilities/state_manager.py

@@ -1,5 +1,6 @@
 import pandas as pd
 import streamlit as st
+from my_model.gen_utilities import free_gpu_resources
 
 
 class StateManager:
@@ -39,6 +40,19 @@ class StateManager:
         df = pd.DataFrame(data)
         st.table(df)
 
+    def load_model(self, detection_model, confidence_level):
+        """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.")
+            free_gpu_resources()
+        except Exception as e:
+            st.error(f"Error loading model: {e}")    
+
     def get_model(self):
         """Retrieve the KBVQA model from the session state."""
         return st.session_state.get('kbvqa', None)