Update app.py

app.py

@@ -11,6 +11,8 @@ from lavis.models import load_model_and_preprocess
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, AutoModelForSeq2SeqLM
 import gradio as gr
 import torch, gc
+from gpuinfo import GPUInfo
+import time
 
 def prepare_data(image, question):
     gc.collect()
@@ -20,6 +22,20 @@ def prepare_data(image, question):
     samples = {"image": image, "text_input": [question]}
     return samples
 
+def running_inf(time_start):
+    time_end = time.time()
+    time_diff = time_end - time_start
+    memory = psutil.virtual_memory()
+    gpu_utilization, gpu_memory = GPUInfo.gpu_usage()
+    gpu_utilization = gpu_utilization[0] if len(gpu_utilization) > 0 else 0
+    gpu_memory = gpu_memory[0] if len(gpu_memory) > 0 else 0
+    system_info = f"""
+    *Memory: {memory.total / (1024 * 1024 * 1024):.2f}GB, used: {memory.percent}%, available: {memory.available / (1024 * 1024 * 1024):.2f}GB.* 
+    *Processing time: {time_diff:.5} seconds.*
+    *GPU Utilization: {gpu_utilization}%, GPU Memory: {gpu_memory}MiB.*
+    """
+    return system_info
+    
 def gradcam_attention(image, question):
     dst_w = 720
     samples = prepare_data(image, question)
@@ -36,11 +52,11 @@ def gradcam_attention(image, question):
     return (avg_gradcam * 255).astype(np.uint8)
 
 def generate_cap(image, question, cap_number):
+    time_start = time.time()
     samples = prepare_data(image, question)
     samples = model.forward_itm(samples=samples)
     samples = model.forward_cap(samples=samples, num_captions=cap_number, num_patches=5)
-    print('Examples of question-guided captions: ')
-    return pd.DataFrame({'Caption': samples['captions'][0][:cap_number]})
+    return pd.DataFrame({'Caption': samples['captions'][0][:cap_number]}), running_inf(time_start)
 
 def postprocess(text):
     for i, ans in enumerate(text):
@@ -51,6 +67,7 @@ def postprocess(text):
     return ans
 
 def generate_answer(image, question):
+    time_start = time.time()
     samples = prepare_data(image, question)
     samples = model.forward_itm(samples=samples)
     samples = model.forward_cap(samples=samples, num_captions=5, num_patches=5)
@@ -67,7 +84,7 @@ def generate_answer(image, question):
     pred_answer = tokenizer.batch_decode(outputs.sequences[:, len(Img2Prompt_input.input_ids[0]):])
     pred_answer = postprocess(pred_answer)
     print(pred_answer, type(pred_answer))
-    return pred_answer
+    return pred_answer, running_inf(time_start)
     
 # setup device to use
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -95,6 +112,7 @@ text_output = gr.Textbox(label="Output Answer")
 demo = gr.Blocks(title=title)
 demo.encrypt = False
 cap_df = gr.DataFrame(value=df_init, label="Caption dataframe", row_count=(0, "dynamic"), max_rows = 20, wrap=True, overflow_row_behaviour='paginate')
+system_info = gr.Markdown(f"*Memory: {memory.total / (1024 * 1024 * 1024):.2f}GB, used: {memory.percent}%, available: {memory.available / (1024 * 1024 * 1024):.2f}GB*")
 
 with demo:
     with gr.Row():
@@ -124,10 +142,10 @@ with demo:
     with gr.Row():
       with gr.Column():
             cap_btn = gr.Button("Generate caption")
-            cap_btn.click(generate_cap, [raw_image, question, number_cap], [cap_df])
+            cap_btn.click(generate_cap, [raw_image, question, number_cap], [cap_df, system_info])
       with gr.Column():
             anws_btn = gr.Button("Answer")
-            anws_btn.click(generate_answer, [raw_image, question], outputs=text_output)
+            anws_btn.click(generate_answer, [raw_image, question], outputs=[text_output, system_info])
     with gr.Row():  
       with gr.Column():  
       #     gradcam_btn = gr.Button("Generate Gradcam")
@@ -135,5 +153,6 @@ with demo:
             cap_df.render()
       with gr.Column():  
             text_output.render()
+            system_info.render()
 
 demo.launch(debug=True)