up

run_demo_multi_models.py

@@ -28,10 +28,24 @@ logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
 device = 0 if torch.cuda.is_available() else "cpu"
-logger.info(f"Model will be loaded on device {device}")
+logger.info(f"Model will be loaded on device `{device}`")
 
 cached_models = {}
 
+
+def print_cuda_memory_info():
+    used_mem, tot_mem = torch.cuda.mem_get_info()
+    logger.info(f"CUDA memory info - Free: {used_mem / 1024 ** 3:.2f} Gb, used: {(tot_mem - used_mem) / 1024 ** 3:.2f} Gb, total: {tot_mem / 1024 ** 3:.2f} Gb")
+
+
+def print_memory_info():
+    # todo
+    if device == "cpu":
+        pass
+    else:
+        print_cuda_memory_info()
+
+
 def maybe_load_cached_pipeline(model_name):
     pipe = cached_models.get(model_name)
     if pipe is None:
@@ -49,6 +63,7 @@ def maybe_load_cached_pipeline(model_name):
         pipe.model.config.max_length = MAX_NEW_TOKENS + 1
 
         logger.info(f"`{model_name}` pipeline has been initialized")
+        print_memory_info()
 
         cached_models[model_name] = pipe
     return pipe
@@ -70,7 +85,7 @@ def transcribe(microphone, file_upload, model_name):
     pipe = maybe_load_cached_pipeline(model_name)
     text = pipe(file)["text"]
 
-    logger.info(f"Transcription: {text}")
+    logger.info(f"Transcription by `{model_name}`: {text}")
 
     return warn_output + text