Update my_model/tabs/run_inference.py

my_model/tabs/run_inference.py

@@ -84,20 +84,24 @@ class InferenceRunner(StateManager):
         st.session_state.button_label = "Reload Model" if self.is_model_loaded() and self.settings_changed else "Load Model"
         
         if st.session_state.method == "Fine-Tuned Model":
-            if self.col1.button(st.session_state.button_label):
-                if st.session_state.button_label == "Load Model":
-                    if self.is_model_loaded():
-                        free_gpu_resources()
-                        self.col1.text("Model already loaded and no settings were changed:)")
+            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":
+                        if self.is_model_loaded():
+                            free_gpu_resources()
+                            self.col1.text("Model already loaded and no settings were changed:)")
+                        else:
+                            free_gpu_resources()
+                            self.load_model()
                     else:
                         free_gpu_resources()
-                        self.load_model()
-                else:
+                        self.reload_detection_model()
+                if nested_col12.button("Force Reload"):
+                    force_reload_model()
+                    
+                if self.is_model_loaded() and st.session_state.kbvqa.all_models_loaded:
                     free_gpu_resources()
-                    self.reload_detection_model()
-
-            if self.is_model_loaded() and st.session_state.kbvqa.all_models_loaded:
-                free_gpu_resources()
-                self.image_qa_app(self.get_model())
+                    self.image_qa_app(self.get_model())
         else:
             self.col1.warning(f'Model using {st.session_state.method} is not deployed yet, will be ready later.')