Update my_model/KBVQA.py

my_model/KBVQA.py

@@ -268,29 +268,30 @@ def prepare_kbvqa_model(only_reload_detection_model: bool = False) -> KBVQA:
     kbvqa = KBVQA()
     kbvqa.detection_model = st.session_state.detection_model
     # Progress bar for model loading
-    with kbvqa.col1.spinner('Loading model...'):
-
-        if not only_reload_detection_model:
-            kbvqa.col1.text('this should take no more than a few minutes!')
-            progress_bar = kbvqa.col1.progress(0)
-            kbvqa.load_detector(kbvqa.detection_model)
-            progress_bar.progress(33)
-            kbvqa.load_caption_model()
-            free_gpu_resources()
-            progress_bar.progress(75)
-            kbvqa.col1.text('Almost there :)')
-            kbvqa.load_fine_tuned_model()
+    with kbvqa.col1:
+        with st.spinner('Loading model...'):
+    
+            if not only_reload_detection_model:
+                st.text('this should take no more than a few minutes!')
+                progress_bar = st.progress(0)
+                kbvqa.load_detector(kbvqa.detection_model)
+                progress_bar.progress(33)
+                kbvqa.load_caption_model()
+                free_gpu_resources()
+                progress_bar.progress(75)
+                st.text('Almost there :)')
+                kbvqa.load_fine_tuned_model()
+                free_gpu_resources()
+                progress_bar.progress(100)
+            else:
+                progress_bar = st.progress(0)
+                kbvqa.load_detector(kbvqa.detection_model)
+                progress_bar.progress(100)
+                
+        if kbvqa.all_models_loaded:
+            st.success('Model loaded successfully and ready for inferecne!')
+            kbvqa.kbvqa_model.eval()
             free_gpu_resources()
-            progress_bar.progress(100)
-        else:
-            progress_bar = kbvqa.col1.progress(0)
-            kbvqa.load_detector(kbvqa.detection_model)
-            progress_bar.progress(100)
-            
-    if kbvqa.all_models_loaded:
-        kbvqa.col1.success('Model loaded successfully and ready for inferecne!')
-        kbvqa.kbvqa_model.eval()
-        free_gpu_resources()
-        return kbvqa
+            return kbvqa