Update my_model/tabs/run_inference.py

my_model/tabs/run_inference.py

@@ -66,6 +66,7 @@ class InferenceRunner(StateManager):
                     qa_history = image_data.get('qa_history', [])
         
                     if image_data['analysis_done']:
+                        st.session_state['loading_in_progress'] = False
                         question = nested_col22.text_input(f"Ask a question about this image ({image_key[-11:]}):", key=f'question_{image_key}')
                         if nested_col22.button('Get Answer', key=f'answer_{image_key}', on_click=self.disable_widgets, disabled=self.is_widget_disabled):
 
@@ -109,10 +110,10 @@ class InferenceRunner(StateManager):
                     if nested_col11.button(st.session_state.button_label, on_click=self.disable_widgets, disabled=self.is_widget_disabled):
                         
                         if st.session_state.button_label == "Load Model":
-                            st.session_state['load_button_clicked'] = True
                             if self.is_model_loaded():
                                 free_gpu_resources()
                                 fine_tuned_model_already_loaded = True
+                      
                             else:
                                 load_fine_tuned_model = True
                         else:
@@ -124,21 +125,26 @@ class InferenceRunner(StateManager):
                 if load_fine_tuned_model:
                     free_gpu_resources()
                     self.load_model()
+                    st.session_state['loading_in_progress'] = False
                     
                 elif fine_tuned_model_already_loaded:
                     free_gpu_resources()
-                    self.col1.text("Model already loaded and no settings were changed:)")      
+                    self.col1.text("Model already loaded and no settings were changed:)")
+                    st.session_state['loading_in_progress'] = False
                     
                 elif reload_detection_model:
                     free_gpu_resources()
                     self.reload_detection_model()
+                    st.session_state['loading_in_progress'] = False
                     
                 elif force_reload_full_model:
                     free_gpu_resources()
                     self.force_reload_model()
+                    st.session_state['loading_in_progress'] = False
                 
             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.')
+                st.session_state['loading_in_progress'] = False
 
         
         if self.is_model_loaded():