Update my_model/tabs/run_inference.py

my_model/tabs/run_inference.py

@@ -86,17 +86,20 @@ class InferenceRunner(StateManager):
                     nested_col21.image(image_for_display, caption=f'Uploaded Image: {image_key[-11:]}')
                     if not image_data['analysis_done'] or self.settings_changed:  # if not done analysis before or even done but settings changed, then we need to analyze again
                         nested_col22.text("Please click 'Analyze Image'..")
+                        free_gpu_resources()
                         with nested_col22:
                             if st.button('Analyze Image', key=f'analyze_{image_key}', on_click=self.disable_widgets, disabled=self.is_widget_disabled):
                                 
                                 caption, detected_objects_str, image_with_boxes = self.analyze_image(image_data['image'])
                                 self.update_image_data(image_key, caption, detected_objects_str, True)
                             st.session_state['loading_in_progress'] = False
+                            free_gpu_resources()
         
                     # Initialize qa_history for each image
                     qa_history = image_data.get('qa_history', [])
                   
                     if image_data['analysis_done']:
+                        free_gpu_resources()
                         st.session_state['loading_in_progress'] = False
                         sample_questions = config.SAMPLE_QUESTIONS.get(image_key, [])
                         selected_question = nested_col22.selectbox(
@@ -118,7 +121,7 @@ class InferenceRunner(StateManager):
                             nested_col22.warning("This question has already been answered.")
                         else:
                             if nested_col22.button('Get Answer', key=f'answer_{image_key}', disabled=self.is_widget_disabled):
-
+                                free_gpu_resources()
                                 answer, prompt_length = self.answer_question(image_data['caption'], image_data['detected_objects_str'], question)
                                 st.session_state['loading_in_progress'] = False
                                 self.add_to_qa_history(image_key, question, answer, prompt_length)
@@ -126,6 +129,7 @@ class InferenceRunner(StateManager):
                     # Display Q&A history and prompts lengths for each image
                     for num, (q, a, p) in enumerate(qa_history):
                         nested_col22.text(f"Q{num+1}: {q}\nA{num+1}: {a}\nPrompt Length: {p}\n")
+                        free_gpu_resources()
 
         
     def run_inference(self):