From d08f739660722780d1db0b87c84d7406ceec6f9f Mon Sep 17 00:00:00 2001
From: Alvaro Bartolome <36760800+alvarobartt@users.noreply.github.com>
Date: Mon, 23 Feb 2026 21:13:55 +0100
Subject: [PATCH] Add `version` display on table and JSON output

---
 src/hf_mem/__init__.py | 3 +++
 src/hf_mem/cli.py      | 6 +++++-
 src/hf_mem/print.py    | 4 ++++
 3 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/src/hf_mem/__init__.py b/src/hf_mem/__init__.py
index e69de29..3f9ac45 100644
--- a/src/hf_mem/__init__.py
+++ b/src/hf_mem/__init__.py
@@ -0,0 +1,3 @@
+from importlib.metadata import version
+
+__version__ = version("hf-mem")
diff --git a/src/hf_mem/cli.py b/src/hf_mem/cli.py
index 11418e0..3f20555 100644
--- a/src/hf_mem/cli.py
+++ b/src/hf_mem/cli.py
@@ -11,6 +11,7 @@
 
 import httpx
 
+from hf_mem import __version__
 from hf_mem.metadata import parse_safetensors_metadata
 from hf_mem.print import print_report
 from hf_mem.types import TorchDtypes, get_safetensors_dtype_bytes, torch_dtype_to_safetensors_dtype
@@ -364,12 +365,15 @@ async def fetch_with_semaphore(url: str) -> Dict[str, Any]:
                     cache_size *= batch_size
 
     if json_output:
-        out = {"model_id": model_id, "revision": revision, **asdict(metadata)}
+        from hf_mem import __version__
+
+        out = {"version": __version__, "model_id": model_id, "revision": revision, **asdict(metadata)}
         if experimental and cache_size:
             out["max_model_len"] = max_model_len
             out["batch_size"] = batch_size
             out["cache_size"] = cache_size
             out["cache_dtype"] = cache_dtype  # type: ignore
+
         print(json.dumps(out))
     else:
         # TODO: Use a `KvCache` dataclass instead and make sure that the JSON output is aligned
diff --git a/src/hf_mem/print.py b/src/hf_mem/print.py
index bcd83f8..1e26b9f 100644
--- a/src/hf_mem/print.py
+++ b/src/hf_mem/print.py
@@ -1,6 +1,7 @@
 import warnings
 from typing import Any, Dict, Literal, Optional
 
+from hf_mem import __version__
 from hf_mem.metadata import SafetensorsMetadata
 
 MIN_NAME_LEN = 5
@@ -179,6 +180,9 @@ def print_report(
         )
     _print_divider(data_col_width + 1, "top")
 
+    _print_row("VERSION", f"hf-mem {__version__}", data_col_width)
+    _print_divider(data_col_width + 1)
+
     if cache:
         total_text = f"{_bytes_to_gib(combined_total):.2f} GiB ({_format_short_number(metadata.param_count)} PARAMS + KV CACHE)"
         total_bar = _make_bar(combined_total, combined_total, data_col_width)