Update my_model/tabs/run_inference.py

my_model/tabs/run_inference.py

@@ -33,7 +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'])
+        st.write('loading in prog?', 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)
@@ -43,7 +43,7 @@ class InferenceRunner(StateManager):
                     self.process_new_image(sample_image_path, image, kbvqa)
 
                     
-        st.write(st.session_state['loading_in_progress'])
+        st.write('loading in prog?',st.session_state['loading_in_progress'])
 
         
         # Image uploader
@@ -118,7 +118,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()
@@ -131,12 +131,13 @@ 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:
+                    t1=time.time()
                     free_gpu_resources()
                     self.load_model()
-                    
+                    st.session_state['time_taken_to_load_model'] = time.time()-t1
                     st.session_state['loading_in_progress'] = False
                     
                 elif fine_tuned_model_already_loaded:
@@ -151,9 +152,9 @@ class InferenceRunner(StateManager):
                     
                 elif force_reload_full_model:
                     free_gpu_resources()
-
+                    t1=time.time()
                     self.force_reload_model()
-
+                    st.session_state['time_taken_to_load_model'] = time.time()-t1
                     st.session_state['loading_in_progress'] = False
                     st.session_state['model_loaded'] = True
                 
@@ -163,11 +164,11 @@ class InferenceRunner(StateManager):
 
         
         if self.is_model_loaded():
-            st.session_state['time_taken_to_load_model'] = time.time()-t1
+            
             st.write(f'time is {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'])
+            st.write('load success', 'loading in prog?', st.session_state['loading_in_progress'])