Update my_model/KBVQA.py

my_model/KBVQA.py

@@ -159,23 +159,29 @@ class KBVQA():
 
         return output_text.capitalize()
 
-def prepare_kbvqa_model(detection_model):
+def prepare_kbvqa_model(detection_model, only_reload_detection_model=False):
     free_gpu_resources()
     kbvqa = KBVQA()
     kbvqa.detection_model = detection_model
     # Progress bar for model loading
     with st.spinner('Loading model...'):
-        
-        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(66)
-        kbvqa.load_fine_tuned_model()
-        free_gpu_resources()
-        progress_bar.progress(100)
+        if not only_reload_detection_model:
+            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(66)
+            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!')