Update my_model/state_manager.py

my_model/state_manager.py

@@ -1,5 +1,6 @@
 import pandas as pd
 import copy
+import time
 from PIL import Image
 import streamlit as st
 from my_model.utilities.gen_utilities import free_gpu_resources
@@ -33,6 +34,8 @@ class StateManager:
             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
+        if 'time_taken_to_load_model' not in st.session_state:
+            st.session_state['time_taken_to_load_model'] = 0
 
             
 
@@ -100,7 +103,7 @@ class StateManager:
         Uses formatted HTML to style the table for better readability.
         """
         self.col3.write("##### Current Model Settings:")
-        data = [{'Setting': key, 'Value': str(value)} for key, value in st.session_state.items() if key in ["confidence_level", 'detection_model', 'method', 'kbvqa', 'previous_state', 'settings_changed', 'loading_in_progress', 'model_loaded' ]]
+        data = [{'Setting': key, 'Value': str(value)} for key, value in st.session_state.items() if key in ["confidence_level", 'detection_model', 'method', 'kbvqa', 'previous_state', 'settings_changed', 'loading_in_progress', 'model_loaded', 'time_taken_to_load_model' ]]
         df = pd.DataFrame(data)
         styled_df = df.style.set_properties(**{'background-color': 'white', 'color': 'black', 'border-color': 'black'}).set_table_styles([{'selector': 'th','props': [('background-color', 'gray'), ('font-weight', 'bold')]}])
         self.col3.write(df)
@@ -129,7 +132,6 @@ class StateManager:
         """
         
         try:
-
             free_gpu_resources()
             st.session_state['kbvqa'] = prepare_kbvqa_model()
             st.session_state['kbvqa'].detection_confidence = st.session_state.confidence_level