Update my_model/KBVQA.py

my_model/KBVQA.py

@@ -50,7 +50,7 @@ class KBVQA:
 
     def __init__(self):
 
-        self.col1, self.col2, self.col3 = st.columns([0.2, 0.6, 0.2])
+       # 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
@@ -245,7 +245,7 @@ class KBVQA:
         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:
-            self.col2.write(f"Prompt too long with {num_tokens} tokens, consider increasing the confidence threshold for the object detector")
+            st.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')
@@ -272,33 +272,33 @@ def prepare_kbvqa_model(only_reload_detection_model: bool = False, force_reload:
     kbvqa = KBVQA()
     kbvqa.detection_model = st.session_state.detection_model
     # Progress bar for model loading
-    with kbvqa.col1:
-        if force_reload:
-            self.delete_model()
-            loading_message = 'Force Reloading model.. this should take no more than a few minutes!'
-        else: loading_message = 'Looading model.. this should take no more than a few minutes!'
-        
-        with st.spinner(loading_message):
-            if not only_reload_detection_model:
-                progress_bar = st.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 :)')
-                kbvqa.load_fine_tuned_model()
-                free_gpu_resources()
-                progress_bar.progress(100)
-            else:
-                progress_bar = st.progress(0)
-                kbvqa.load_detector(kbvqa.detection_model)
-                progress_bar.progress(100)
-                
-        if kbvqa.all_models_loaded:
-            st.success('Model loaded successfully and ready for inferecne!')
-            kbvqa.kbvqa_model.eval()
+  
+    if force_reload:
+        self.delete_model()
+        loading_message = 'Force Reloading model.. this should take no more than a few minutes!'
+    else: loading_message = 'Looading model.. this should take no more than a few minutes!'
+    
+    with st.spinner(loading_message):
+        if not only_reload_detection_model:
+            progress_bar = st.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 :)')
+            kbvqa.load_fine_tuned_model()
             free_gpu_resources()
-            return kbvqa
+            progress_bar.progress(100)
+        else:
+            progress_bar = st.progress(0)
+            kbvqa.load_detector(kbvqa.detection_model)
+            progress_bar.progress(100)
+            
+    if kbvqa.all_models_loaded:
+        st.success('Model loaded successfully and ready for inferecne!')
+        kbvqa.kbvqa_model.eval()
+        free_gpu_resources()
+        return kbvqa