Update my_model/tabs/run_inference.py

my_model/tabs/run_inference.py

@@ -73,7 +73,10 @@ class InferenceRunner(StateManager):
                     for q, a in qa_history:
                         nested_col22.text(f"Q: {q}\nA: {a}\n")
 
+    def display_message(self, message, warning=False, write=False, text=False):
+        pass
 
+        
     def run_inference(self):
         self.set_up_widgets()
         st.session_state['settings_changed'] = self.has_state_changed()
@@ -89,21 +92,37 @@ 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 st.session_state.button_label == "Load Model":
+                            st.session_state['load_button_clicked'] = True
                             if self.is_model_loaded():
                                 free_gpu_resources()
-                                self.col1.text("Model already loaded and no settings were changed:)")
+                                fine_tuned_model_already_loaded = True
                             else:
-                                free_gpu_resources()
-                                self.load_model()
+                                load_fine_tuned_model = True
                         else:
-                            free_gpu_resources()
-                            self.reload_detection_model()
+                            reload_detection_model = True
+                            
                     if nested_col12.button("Force Reload"):
-                        self.force_reload_model()
+                        force_reload_full_model = True
                         
                     if self.is_model_loaded() and st.session_state.kbvqa.all_models_loaded:
                         free_gpu_resources()
                         self.image_qa_app(self.get_model())
-            else:
+
+            elif load_fine_tuned_model:
+                free_gpu_resources()
+                self.load_model()
+                
+            elif fine_tuned_model_already_loaded:
+                self.col1.text("Model already loaded and no settings were changed:)")      
+                
+            elif reload_detection_model:
+                free_gpu_resources()
+                
+                self.reload_detection_model()
+            elif force_reload_full_model:
+                self.force_reload_model()
+                
+            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.')