Update my_model/KBVQA.py

my_model/KBVQA.py

@@ -176,12 +176,16 @@ class KBVQA:
         free_gpu_resources()
         if self.kbvqa_model is not None:
             del self.kbvqa_model
+            free_gpu_resources()
         if self.captioner is not None:
             del self.captioner
+            free_gpu_resources()
         if self.detector is not None:
             del self.detector
-
+            free_gpu_resources()
+            
         free_gpu_resources()
+        prepare_kbvqa_model(only_reload_detection_model=False, force_reload=True)
 
 
     def format_prompt(self, current_query: str, history: Optional[str] = None, sys_prompt: Optional[str] = None, caption: str = None, objects: Optional[str] = None) -> str:
@@ -253,7 +257,7 @@ class KBVQA:
 
         return output_text.capitalize()
 
-def prepare_kbvqa_model(only_reload_detection_model: bool = False) -> KBVQA:
+def prepare_kbvqa_model(only_reload_detection_model: bool = False, force_reload: bool = True) -> KBVQA:
     """
     Prepares the KBVQA model for use, including loading necessary sub-models.
 
@@ -269,7 +273,11 @@ def prepare_kbvqa_model(only_reload_detection_model: bool = False) -> KBVQA:
     kbvqa.detection_model = st.session_state.detection_model
     # Progress bar for model loading
     with kbvqa.col1:
-        with st.spinner('Loading model.. this should take no more than a few minutes!'):
+        if force_reload:
+            loading_message = 'Force Reloading model.. this should take no more than a few minutes!'
+        else: loading_message = 'Looading model.. this should take no more than a few minutes!'
+        
+        with st.spinner(loading_message):
             if not only_reload_detection_model:
                 progress_bar = st.progress(0)
                 kbvqa.load_detector(kbvqa.detection_model)