Update my_model/KBVQA.py

my_model/KBVQA.py

@@ -21,6 +21,7 @@ class KBVQA():
         self.kbvqa_tokenizer = None
         self.captioner = None
         self.detector = None
+        sel.detection_model = None
         self.kbvqa_model = None
         self.access_token = os.getenv("HUGGINGFACE_TOKEN")
       #  self.kbvqa_model_loaded = self.all_models_loaded() 
@@ -87,6 +88,22 @@ class KBVQA():
     def all_models_loaded(self):
         return self.kbvqa_model is not None and self.captioner is not None and self.detector is not None
 
+    def force_reload_model(self):
+        free_gpu_resources()
+        if self.kbvqa_model is not None:
+            del self.kbvqa_model
+        if self.captioner is not None:
+            del self.captioner
+        if self.detector is not None:
+            del self.detector
+
+        free_gpu_resources()
+
+        
+
+        
+            
+
 
 
     def format_prompt(self, current_query, history = None , sys_prompt=None, caption=None, objects=None):
@@ -144,11 +161,13 @@ class KBVQA():
 def prepare_kbvqa_model(detection_model):
     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(detection_model)
+
+        kbvqa.load_detector(kbvqa.detection_model)
         progress_bar.progress(33)
         kbvqa.load_caption_model()
         free_gpu_resources()