Update my_model/KBVQA.py

my_model/KBVQA.py

@@ -11,7 +11,7 @@ from my_model.object_detection import ObjectDetector
 import my_model.config.kbvqa_config as config
 
 
-class KBVQA():
+class KBVQA(StateManager):
     """
     The KBVQA class encapsulates the functionality for the Knowledge-Based Visual Question Answering (KBVQA) model. 
     It integrates various components such as an image captioning model, object detection model, and a fine-tuned 
@@ -48,7 +48,8 @@ class KBVQA():
     """
 
     def __init__(self):
- 
+
+        super().__init__()
         self.kbvqa_model_name: str = config.KBVQA_MODEL_NAME
         self.quantization: str = config.QUANTIZATION
         self.max_context_window: int = config.MAX_CONTEXT_WINDOW
@@ -269,14 +270,14 @@ def prepare_kbvqa_model(only_reload_detection_model: bool = False) -> KBVQA:
     with st.spinner('Loading model...'):
 
         if not only_reload_detection_model:
-            st.text('this should take no more than a few minutes!')
-            progress_bar = st.progress(0)
+            self.col1.text('this should take no more than a few minutes!')
+            progress_bar = self.col1.progress(0)
             kbvqa.load_detector(kbvqa.detection_model)
             progress_bar.progress(33)
             kbvqa.load_caption_model()
             free_gpu_resources()
             progress_bar.progress(75)
-            st.text('Almost there :)')
+            self.col1.text('Almost there :)')
             kbvqa.load_fine_tuned_model()
             free_gpu_resources()
             progress_bar.progress(100)
@@ -286,7 +287,7 @@ def prepare_kbvqa_model(only_reload_detection_model: bool = False) -> KBVQA:
             progress_bar.progress(100)
             
     if kbvqa.all_models_loaded:
-        st.success('Model loaded successfully and ready for inferecne!')
+        self.col1.success('Model loaded successfully and ready for inferecne!')
         kbvqa.kbvqa_model.eval()
         free_gpu_resources()
         return kbvqa