Update my_model/tabs/run_inference.py

my_model/tabs/run_inference.py

@@ -4,6 +4,7 @@ import bitsandbytes
 import accelerate
 import scipy
 import copy
+import time
 from PIL import Image
 import torch.nn as nn
 import pandas as pd
@@ -32,6 +33,7 @@ class InferenceRunner(StateManager):
         # Display sample images as clickable thumbnails
         self.col1.write("Choose from sample images:")
         cols = self.col1.columns(len(self.sample_images))
+        st.write(st.session_state['loading_in_progress'])
         for idx, sample_image_path in enumerate(self.sample_images):
             with cols[idx]:
                 image = Image.open(sample_image_path)
@@ -108,7 +110,7 @@ class InferenceRunner(StateManager):
                 with st.container():
                     nested_col11, nested_col12 = st.columns([0.5, 0.5])
                     if nested_col11.button(st.session_state.button_label, on_click=self.disable_widgets, disabled=self.is_widget_disabled):
-                        
+                        t1=time.time()
                         if st.session_state.button_label == "Load Model":
                             if self.is_model_loaded():
                                 free_gpu_resources()
@@ -121,10 +123,12 @@ class InferenceRunner(StateManager):
                             
                     if nested_col12.button("Force Reload", on_click=self.disable_widgets, disabled=self.is_widget_disabled):
                         force_reload_full_model = True
+                        t1=time.time()
 
                 if load_fine_tuned_model:
                     free_gpu_resources()
                     self.load_model()
+                    
                     st.session_state['loading_in_progress'] = False
                     
                 elif fine_tuned_model_already_loaded:
@@ -139,8 +143,11 @@ class InferenceRunner(StateManager):
                     
                 elif force_reload_full_model:
                     free_gpu_resources()
+
                     self.force_reload_model()
+
                     st.session_state['loading_in_progress'] = False
+                    st.session_state['model_loaded'] = True
                 
             elif st.session_state.method == "In-Context Learning (n-shots)":
                 self.col1.warning(f'Model using {st.session_state.method} is not deployed yet, will be ready later.')
@@ -148,8 +155,9 @@ class InferenceRunner(StateManager):
 
         
         if self.is_model_loaded():
-            st.session_state['loading_in_progress'] = False
+            st.session_state['time_taken_to_load_model'] = time.time()-t1
             free_gpu_resources()
+            st.session_state['loading_in_progress'] = False
             self.image_qa_app(self.get_model())
             st.write(st.session_state['loading_in_progress'])