Update my_model/KBVQA.py

my_model/KBVQA.py

@@ -9,7 +9,7 @@ from my_model.utilities.gen_utilities import free_gpu_resources
 from my_model.captioner.image_captioning import ImageCaptioningModel
 from my_model.object_detection import ObjectDetector
 import my_model.config.kbvqa_config as config
-from my_model.state_manager import StateManager
+
 
 
 class KBVQA(StateManager):
@@ -50,7 +50,7 @@ class KBVQA(StateManager):
 
     def __init__(self):
 
-        super().__init__()
+        self.col1, self.col2, self.col3 = st.columns([0.2, 0.6, 0.2])
         self.kbvqa_model_name: str = config.KBVQA_MODEL_NAME
         self.quantization: str = config.QUANTIZATION
         self.max_context_window: int = config.MAX_CONTEXT_WINDOW
@@ -241,7 +241,7 @@ class KBVQA(StateManager):
         prompt = self.format_prompt(question, caption=caption, objects=detected_objects_str)
         num_tokens = len(self.kbvqa_tokenizer.tokenize(prompt))
         if num_tokens > self.max_context_window:
-            st.write(f"Prompt too long with {num_tokens} tokens, consider increasing the confidence threshold for the object detector")
+            self.col2.write(f"Prompt too long with {num_tokens} tokens, consider increasing the confidence threshold for the object detector")
             return
 
         model_inputs = self.kbvqa_tokenizer(prompt, add_special_tokens=False, return_tensors="pt").to('cuda')
@@ -268,7 +268,7 @@ def prepare_kbvqa_model(only_reload_detection_model: bool = False) -> KBVQA:
     kbvqa = KBVQA()
     kbvqa.detection_model = st.session_state.detection_model
     # Progress bar for model loading
-    with st.spinner('Loading model...'):
+    with self.col1.spinner('Loading model...'):
 
         if not only_reload_detection_model:
             self.col1.text('this should take no more than a few minutes!')
@@ -283,7 +283,7 @@ def prepare_kbvqa_model(only_reload_detection_model: bool = False) -> KBVQA:
             free_gpu_resources()
             progress_bar.progress(100)
         else:
-            progress_bar = st.progress(0)
+            progress_bar = self.col1.progress(0)
             kbvqa.load_detector(kbvqa.detection_model)
             progress_bar.progress(100)