Update app.py

show system info

app.py

@@ -13,6 +13,11 @@ from mediapipe.framework.formats import landmark_pb2
 from mediapipe.tasks import python
 from mediapipe.tasks.python import vision
 import cv2
+import psutil
+from gpuinfo import GPUInfo
+import time
+import gc
+import torch
 
 from diffusers import (
     FlaxControlNetModel,
@@ -210,7 +215,22 @@ def infer(prompt, negative_prompt, image, model_type="Standard"):
     images = images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:])
 
     results = [i for i in images]
-    return [overlap_image, annotated_image] + results
+    
+    # running info
+    time_end = time.time()
+    time_diff = time_end - time_start
+    gc.collect()
+    torch.cuda.empty_cache()
+    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 [overlap_image, annotated_image] + results, system_info 
 
 
 with gr.Blocks(theme='gradio/soft') as demo:
@@ -301,7 +321,7 @@ with gr.Blocks(theme='gradio/soft') as demo:
             submit_btn = gr.Button(value = "Submit")
             # inputs = [prompt_input, negative_prompt, input_image]
             # submit_btn.click(fn=infer, inputs=inputs, outputs=[output_image])
-    
+            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 gr.Column():
             output_image = gr.Gallery(label='Output Image', show_label=False, elem_id="gallery").style(grid=2, height='auto')
     
@@ -334,12 +354,12 @@ with gr.Blocks(theme='gradio/soft') as demo:
             ],
         ],
         inputs=[prompt_input, negative_prompt, input_image, model_type],
-        outputs=[output_image],
+        outputs=[output_image, system_info],
         fn=infer,
         cache_examples=True,
     )
             
     inputs = [prompt_input, negative_prompt, input_image, model_type]
-    submit_btn.click(fn=infer, inputs=inputs, outputs=[output_image])
+    submit_btn.click(fn=infer, inputs=inputs, outputs=[output_image, system_info])
 
 demo.launch()