Update my_model/tabs/run_inference.py

my_model/tabs/run_inference.py

@@ -56,8 +56,10 @@ class InferenceRunner(StateManager):
                         nested_col22.text("Please click 'Analyze Image'..")
                         with nested_col22:
                             if st.button('Analyze Image', key=f'analyze_{image_key}'):
+                                st.session_state['loading_in_progress'] = True
                                 caption, detected_objects_str, image_with_boxes = self.analyze_image(image_data['image'], kbvqa)
                                 self.update_image_data(image_key, caption, detected_objects_str, True)
+                            st.session_state['loading_in_progress'] = False
         
                     # Initialize qa_history for each image
                     qa_history = image_data.get('qa_history', [])
@@ -65,10 +67,13 @@ class InferenceRunner(StateManager):
                     if image_data['analysis_done']:
                         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}'):
+                            st.session_state['loading_in_progress'] = True
                             if question not in [q for q, _ in qa_history]:
                                 answer = self.answer_question(image_data['caption'], image_data['detected_objects_str'], question, kbvqa)
                                 self.add_to_qa_history(image_key, question, answer)
                             else: nested_col22.warning("This questions has already been answered.")
+                                
+                            st.session_state['loading_in_progress'] = False
         
                     # Display Q&A history for each image
                     for q, a in qa_history:
@@ -99,7 +104,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):
+                    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
@@ -111,7 +116,7 @@ class InferenceRunner(StateManager):
                         else:
                             reload_detection_model = True
                             
-                    if nested_col12.button("Force Reload"):
+                    if nested_col12.button("Force Reload", on_click=self.disable_widgets, disabled=self.is_widget_disabled):
                         force_reload_full_model = True
 
                 if load_fine_tuned_model: