add memory usage

requirements.txt

@@ -4,4 +4,5 @@ Ninja
 gradio
 huggingface_hub
 hf_transfer
-Pillow==9.5.0
+Pillow==9.5.0
+psutil

visualizer_drag_gradio.py

@@ -7,6 +7,8 @@ from functools import partial
 from pathlib import Path
 import time
 
+import psutil
+
 import gradio as gr
 import numpy as np
 import torch
@@ -160,6 +162,24 @@ def preprocess_mask_info(global_state, image):
     return global_state
 
 
+def print_memory_usage():
+    # Print system memory usage
+    print(f"System memory usage: {psutil.virtual_memory().percent}%")
+
+    # Print GPU memory usage
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+        print(f"GPU memory usage: {torch.cuda.memory_allocated() / 1e9} GB")
+        print(
+            f"Max GPU memory usage: {torch.cuda.max_memory_allocated() / 1e9} GB")
+        device_properties = torch.cuda.get_device_properties(device)
+        available_memory = device_properties.total_memory - \
+            torch.cuda.max_memory_allocated()
+        print(f"Available GPU memory: {available_memory / 1e9} GB")
+    else:
+        print("No GPU available")
+
+
 # filter large models running on SPACES
 allowed_checkpoints = []  # all checkpoints
 if IS_SPACE:
@@ -557,6 +577,7 @@ with gr.Blocks() as app:
             step_idx = 0
             last_time = time.time()
             while True:
+                print_memory_usage()
                 # add a TIMEOUT break
                 print(f'Running time: {time.time() - last_time}')
                 if IS_SPACE and time.time() - last_time > TIMEOUT: