Update my_model/tabs/run_inference.py

my_model/tabs/run_inference.py

@@ -82,26 +82,28 @@ 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":
-            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:)")
+
+        with self.col1:
+            if st.session_state.method == "Fine-Tuned Model":
+                
+                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"):
+                        self.force_reload_model()
+                        
+                    if self.is_model_loaded() and st.session_state.kbvqa.all_models_loaded:
                         free_gpu_resources()
-                        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.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.')
+                        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.')