Update my_model/tabs/run_inference.py

my_model/tabs/run_inference.py

@@ -94,6 +94,11 @@ class InferenceRunner(StateManager):
         st.session_state.button_label = "Reload Model" if self.is_model_loaded() and self.settings_changed else "Load Model"
 
         with self.col1:
+            
+            if self.is_model_loaded():
+                free_gpu_resources()
+                self.image_qa_app(self.get_model())
+                
             if st.session_state.method == "Fine-Tuned Model":
                 
                 with st.container():
@@ -112,10 +117,6 @@ class InferenceRunner(StateManager):
                             
                     if nested_col12.button("Force Reload"):
                         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())
 
                 if load_fine_tuned_model:
                     free_gpu_resources()
@@ -131,6 +132,8 @@ class InferenceRunner(StateManager):
                 elif force_reload_full_model:
                     free_gpu_resources()
                     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.')