Update my_model/KBVQA.py

my_model/KBVQA.py

@@ -99,6 +99,7 @@ class KBVQA:
         
         self.captioner = ImageCaptioningModel()
         self.captioner.load_model()
+        free_gpu_resources()
 
     def get_caption(self, img: Image.Image) -> str:
         """
@@ -110,8 +111,9 @@ class KBVQA:
         Returns:
             str: The generated caption for the image.
         """
-
-        return self.captioner.generate_caption(img)
+        caption = self.captioner.generate_caption(img)
+        free_gpu_resources()
+        return caption
 
     def load_detector(self, model: str) -> None:
         """
@@ -123,6 +125,7 @@ class KBVQA:
 
         self.detector = ObjectDetector()
         self.detector.load_model(model)
+        free_gpu_resources()
 
     def detect_objects(self, img: Image.Image) -> Tuple[Image.Image, str]:
         """
@@ -136,8 +139,11 @@ class KBVQA:
         """
       
         image = self.detector.process_image(img)
+        free_gpu_resources()
         detected_objects_string, detected_objects_list = self.detector.detect_objects(image, threshold=st.session_state['confidence_level'])
+        free_gpu_resources()
         image_with_boxes = self.detector.draw_boxes(img, detected_objects_list)
+        free_gpu_resources()
         return image_with_boxes, detected_objects_string
 
     def load_fine_tuned_model(self) -> None:
@@ -150,6 +156,8 @@ class KBVQA:
                                                                 low_cpu_mem_usage=True, 
                                                                 quantization_config=self.bnb_config,
                                                                 token=self.access_token)
+
+        free_gpu_resources()
         
         self.kbvqa_tokenizer = AutoTokenizer.from_pretrained(self.kbvqa_model_name, 
                                                              use_fast=self.use_fast, 
@@ -157,7 +165,7 @@ class KBVQA:
                                                              trust_remote_code=self.trust_remote, 
                                                              add_eos_token=self.add_eos_token,
                                                              token=self.access_token)
-
+        free_gpu_resources()
 
     @property
     def all_models_loaded(self):
@@ -225,7 +233,7 @@ class KBVQA:
         Returns:
             str: The generated answer to the question.
         """
-        
+        free_gpu_resources()
         prompt = self.format_prompt(question, caption=caption, objects=detected_objects_str)
         num_tokens = len(self.kbvqa_tokenizer.tokenize(prompt))
         self.current_prompt_length = num_tokens
@@ -234,8 +242,10 @@ class KBVQA:
             return
 
         model_inputs = self.kbvqa_tokenizer(prompt, add_special_tokens=False, return_tensors="pt").to('cuda')
+        free_gpu_resources()
         input_ids = model_inputs["input_ids"]
         output_ids = self.kbvqa_model.generate(input_ids)
+        free_gpu_resources()
         index = input_ids.shape[1] # needed to avoid printing the input prompt
         history = self.kbvqa_tokenizer.decode(output_ids[0], skip_special_tokens=False)
         output_text = self.kbvqa_tokenizer.decode(output_ids[0][index:], skip_special_tokens=True)