up

run_demo_openai.py

@@ -3,6 +3,7 @@ import warnings
 
 import gradio as gr
 import pytube as pt
+import psutil
 import torch
 import whisper
 from huggingface_hub import hf_hub_download, model_info
@@ -42,6 +43,13 @@ logger.info(f"Model will be loaded on device `{device}`")
 cached_models = {}
 
 
+def _print_memory_info():
+    memory = psutil.virtual_memory()
+    logger.info(
+        f"Memory: {memory.total / (1024 ** 3):.2f}GB, used: {memory.percent}%, available: {memory.available / (1024 ** 3):.2f}GB"
+    )
+
+
 def print_cuda_memory_info():
     used_mem, tot_mem = torch.cuda.mem_get_info()
     logger.info(
@@ -50,11 +58,8 @@ def print_cuda_memory_info():
 
 
 def print_memory_info():
-    # todo
-    if device == "cpu":
-        pass
-    else:
-        print_cuda_memory_info()
+    _print_memory_info()
+    print_cuda_memory_info()
 
 
 def maybe_load_cached_pipeline(model_name):