table memory

infero/main.py

@@ -11,6 +11,7 @@
     print_neutral,
     print_success_bold,
     print_error,
+    get_memory_usage,
 )
 from infero.pull.models import remove_model
 
@@ -19,6 +20,7 @@
 
 @app.command("run")
 def run(model: str, quantize: bool = False):
+    print_neutral(f"{get_memory_usage()/ 1024 / 1024} MB")
     if check_model(model):
         model_path = os.path.join(get_models_dir(), sanitize_model_name(model))
         package_dir = get_package_dir()
@@ -50,12 +52,19 @@ def list_models():
     models = []
     for model in os.listdir(models_dir):
         quantized = (
-            "✅"
+            f"{os.path.getsize(os.path.join(models_dir, model, 'model_quantized.onnx')) / 1024 / 1024:.2f}"
             if os.path.exists(os.path.join(models_dir, model, "model_quantized.onnx"))
             else ""
         )
-        models.append([model, quantized])
-    table = tabulate(models, headers=["Name", "Quantized"], tablefmt="grid")
+        size = (
+            os.path.getsize(os.path.join(models_dir, model, "pytorch_model.bin"))
+            / 1024
+            / 1024
+        )
+        models.append([model, size, quantized])
+    table = tabulate(
+        models, headers=["Name", "Size (MB)", "Quantized (MB)"], tablefmt="grid"
+    )
     print_neutral(table)