Update my_model/KBVQA.py

my_model/KBVQA.py

@@ -3,6 +3,7 @@ import torch
 import os
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from typing import Optional
+from my_model.utilities import free_gpu_resources
 from my_model.captioner.image_captioning import ImageCaptioningModel
 from my_model.object_detection import ObjectDetector
 
@@ -141,20 +142,25 @@ class KBVQA():
         return output_text.capitalize()
 
 def prepare_kbvqa_model(detection_model):
+    free_gpu_resources()
     kbvqa = KBVQA()
     # Progress bar for model loading
     with st.spinner('Loading models...'):
+        
         progress_bar = st.progress(0)
         kbvqa.load_detector(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 kbvqa.all_models_loaded:
         st.success('Model loaded successfully!')
         kbvqa.kbvqa_model.eval()
+        free_gpu_resources()
         return kbvqa