hf-mem 0.0.3__py3-none-any.whl

hf_mem/cli.py

@@ -0,0 +1,195 @@
+import argparse
+import asyncio
+import json
+import math
+import os
+import struct
+import sys
+from functools import reduce
+from typing import Any, Dict, List, Optional
+
+import httpx
+
+from hf_mem.print import print_report_for_diffusers, print_report_for_transformers
+
+# NOTE: Defines the bytes that will be fetched per safetensors file, but the metadata
+# can indeed be larger than that
+MAX_METADATA_SIZE = 100_000
+REQUEST_TIMEOUT = float(os.getenv("REQUEST_TIMEOUT", 10.0))
+MAX_CONCURRENCY = int(os.getenv("MAX_WORKERS", min(32, (os.cpu_count() or 1) + 4)))
+
+
+# NOTE: Return type-hint set to `Any`, but it will only be a JSON-compatible object
+async def get_json_file(
+    client: httpx.AsyncClient, url: str, headers: Optional[Dict[str, str]] = None
+) -> Any:
+    response = await client.get(url, headers=headers, timeout=REQUEST_TIMEOUT)
+    response.raise_for_status()
+    return response.json()
+
+
+async def fetch_safetensors_metadata(
+    client: httpx.AsyncClient, url: str, headers: Optional[Dict[str, str]] = None
+) -> Dict[str, Any]:
+    headers = {"Range": f"bytes=0-{MAX_METADATA_SIZE}", **(headers or {})}
+    response = await client.get(url, headers=headers, timeout=REQUEST_TIMEOUT)
+    response.raise_for_status()
+
+    metadata = response.read()
+    # NOTE: Parse the first 8 bytes as a little-endian uint64 (size of the metadata)
+    metadata_size = struct.unpack("<Q", metadata[:8])[0]
+
+    if metadata_size < MAX_METADATA_SIZE:
+        metadata = metadata[8 : metadata_size + 8]
+        return json.loads(metadata)
+
+    # NOTE: Given that by default we just fetch the first 100_000 bytes, if the content is larger
+    # then we simply fetch the remainder again
+    metadata = metadata[8 : MAX_METADATA_SIZE + 8]
+    headers["Range"] = f"bytes={MAX_METADATA_SIZE + 1}-{metadata_size + 7}"
+
+    response = await client.get(url, headers=headers, timeout=REQUEST_TIMEOUT)
+    response.raise_for_status()
+
+    metadata += response.read()
+    return json.loads(metadata)
+
+
+async def run(model_id: str, revision: str) -> None:
+    headers = {}
+    if token := os.getenv("HF_TOKEN", None):
+        headers["Authorization"] = f"Bearer {token}"
+
+    client = httpx.AsyncClient(
+        limits=httpx.Limits(
+            max_keepalive_connections=MAX_CONCURRENCY,
+            max_connections=MAX_CONCURRENCY,
+        ),
+        timeout=httpx.Timeout(REQUEST_TIMEOUT),
+        http2=True,
+        follow_redirects=True,
+    )
+
+    # TODO: `recursive=true` shouldn't really be required unless it's a Diffusers
+    # models... I don't think this adds extra latency anyway
+    url = f"https://huggingface.co/api/models/{model_id}/tree/{revision}?recursive=true"
+
+    files = await get_json_file(client=client, url=url, headers=headers)
+
+    file_paths = [
+        f["path"]
+        for f in files
+        if f.get("path", None) is not None and f.get("type", None) == "file"
+    ]
+
+    if "model.safetensors" in file_paths:
+        url = f"https://huggingface.co/{model_id}/resolve/{revision}/model.safetensors"
+        metadata = await fetch_safetensors_metadata(
+            client=client, url=url, headers=headers
+        )
+        print_report_for_transformers(model_id, revision, metadata)
+    elif "model.safetensors.index.json" in file_paths:
+        # TODO: We could eventually skip this request in favour of a greedy approach on trying to pull all the
+        # files following the formatting `model-00000-of-00000.safetensors`
+        url = f"https://huggingface.co/{model_id}/resolve/{revision}/model.safetensors.index.json"
+        files_index = await get_json_file(client=client, url=url, headers=headers)
+
+        urls = {
+            f"https://huggingface.co/{model_id}/resolve/{revision}/{f}"
+            for f in set(files_index["weight_map"].values())
+        }
+
+        semaphore = asyncio.Semaphore(MAX_CONCURRENCY)
+
+        async def fetch_semaphore(url: str) -> Dict[str, Any]:
+            async with semaphore:
+                return await fetch_safetensors_metadata(
+                    client=client, url=url, headers=headers
+                )
+
+        tasks = [asyncio.create_task(fetch_semaphore(url)) for url in urls]
+        metadata_list: List[Dict[str, Any]] = await asyncio.gather(
+            *tasks, return_exceptions=False
+        )
+
+        metadata = reduce(lambda acc, metadata: acc | metadata, metadata_list, {})
+        print_report_for_transformers(model_id, revision, metadata)
+    elif "model_index.json" in file_paths:
+        url = f"https://huggingface.co/{model_id}/resolve/{revision}/model_index.json"
+        files_index = await get_json_file(client=client, url=url, headers=headers)
+        paths = {k for k, _ in files_index.items() if not k.startswith("_")}
+
+        path_urls: Dict[str, List[str]] = {}
+        for path in paths:
+            if f"{path}/diffusion_pytorch_model.safetensors" in file_paths:
+                path_urls[path] = [
+                    f"https://huggingface.co/{model_id}/resolve/{revision}/{path}/diffusion_pytorch_model.safetensors"
+                ]
+            elif f"{path}/model.safetensors" in file_paths:
+                path_urls[path] = [
+                    f"https://huggingface.co/{model_id}/resolve/{revision}/{path}/model.safetensors"
+                ]
+            elif f"{path}/diffusion_pytorch_model.safetensors.index.json" in file_paths:
+                url = f"https://huggingface.co/{model_id}/resolve/{revision}/{path}/diffusion_pytorch_model.safetensors.index.json"
+                files_index = await get_json_file(
+                    client=client, url=url, headers=headers
+                )
+                path_urls[path] = list({
+                    f"https://huggingface.co/{model_id}/resolve/{revision}/{path}/{f}"
+                    for _, f in files_index["weight_map"].items()  # type: ignore
+                })
+            elif f"{path}/model.safetensors.index.json" in file_paths:
+                url = f"https://huggingface.co/{model_id}/resolve/{revision}/{path}/model.safetensors.index.json"
+                files_index = await get_json_file(
+                    client=client, url=url, headers=headers
+                )
+                path_urls[path] = list({
+                    f"https://huggingface.co/{model_id}/resolve/{revision}/{path}/{f}"
+                    for _, f in files_index["weight_map"].items()  # type: ignore
+                })
+
+        semaphore = asyncio.Semaphore(MAX_CONCURRENCY)
+
+        async def fetch_semaphore(url: str) -> Dict[str, Any]:
+            async with semaphore:
+                return await fetch_safetensors_metadata(
+                    client=client, url=url, headers=headers
+                )
+
+        metadata = {}
+        for path, urls in path_urls.items():
+            tasks = [asyncio.create_task(fetch_semaphore(url)) for url in urls]
+            metadata_list: List[Dict[str, Any]] = await asyncio.gather(
+                *tasks, return_exceptions=False
+            )
+            metadata[path] = reduce(
+                lambda acc, metadata: acc | metadata, metadata_list, {}
+            )
+
+        print_report_for_diffusers(
+            model_id=model_id, revision=revision, metadata=metadata
+        )
+    else:
+        raise RuntimeError(
+            "NONE OF `model.safetensors`, `model.safetensors.index.json`, `model_index.json` HAS BEEN FOUND"
+        )
+
+    sys.exit(0)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model-id", required=True, help="Model ID on the Hugging Face Hub"
+    )
+    parser.add_argument(
+        "--revision",
+        default="main",
+        help="Model revision on the Hugging Face Hub",
+    )
+
+    args = parser.parse_args()
+    model_id = args.model_id
+    revision = args.revision
+
+    asyncio.run(run(model_id, revision))

hf_mem/print.py

@@ -0,0 +1,290 @@
+import math
+from typing import Any, Dict, Literal, Optional
+
+MIN_NAME_LEN = 5
+MAX_NAME_LEN = 13
+MIN_DATA_LEN = 20
+MAX_DATA_LEN = 32
+BORDERS_AND_PADDING = 7
+
+BOX = {
+    "tl": "┌",
+    "tr": "┐",
+    "bl": "└",
+    "br": "┘",
+    "ht": "─",
+    "vt": "│",
+    "tsep": "┬",
+    "bsep": "┴",
+    "lm": "├",
+    "rm": "┤",
+    "mm": "┼",
+}
+
+
+def _print_with_color(content: str) -> None:
+    print(f"\x1b[38;2;244;183;63m{content}\x1b[0m")
+
+
+def _print_header(current_len: int) -> None:
+    length = current_len + MAX_NAME_LEN + BORDERS_AND_PADDING
+    top = BOX["tl"] + (BOX["tsep"] * (length - 2)) + BOX["tr"]
+    _print_with_color(top)
+
+    bottom = BOX["lm"] + (BOX["bsep"] * (length - 2)) + BOX["rm"]
+    _print_with_color(bottom)
+
+
+def _print_centered(text: str, current_len: int) -> None:
+    max_len = current_len + MAX_NAME_LEN - BORDERS_AND_PADDING
+    total_width = max_len + 12
+    text_len = len(text)
+    pad_left = (total_width - text_len) // 2
+    pad_right = total_width - text_len - pad_left
+    _print_with_color(f"{BOX['vt']}{' ' * pad_left}{text}{' ' * pad_right}{BOX['vt']}")
+
+
+def _print_divider(
+    current_len: int,
+    side: Optional[Literal["top", "top-continue", "bottom", "bottom-continue"]] = None,
+) -> None:
+    match side:
+        case "top":
+            left, mid, right = BOX["lm"], BOX["tsep"], BOX["rm"]
+        case "top-continue":
+            left, mid, right = BOX["lm"], BOX["bsep"], BOX["rm"]
+        case "bottom":
+            left, mid, right = BOX["bl"], BOX["bsep"], BOX["br"]
+        case "bottom-continue":
+            left, mid, right = BOX["lm"], BOX["bsep"], BOX["rm"]
+        case _:
+            left, mid, right = BOX["lm"], BOX["mm"], BOX["rm"]
+
+    name_col_inner = MAX_NAME_LEN + 2
+    data_col_inner = current_len + 1
+
+    line = left
+    line += BOX["ht"] * name_col_inner
+    line += mid
+    line += BOX["ht"] * data_col_inner
+    line += right
+    _print_with_color(line)
+
+
+def _format_name(name: str) -> str:
+    if len(name) < MIN_NAME_LEN:
+        return f"{name:<{MIN_NAME_LEN}}"
+    if len(name) > MAX_NAME_LEN:
+        return name[: MAX_NAME_LEN - 3] + "..."
+    return f"{name:<{MAX_NAME_LEN}}"
+
+
+def _print_row(name: str, text: str, current_len: int) -> None:
+    name_fmt = _format_name(name)
+    data_fmt = f"{str(text):<{current_len}}"
+    _print_with_color(f"{BOX['vt']} {name_fmt} {BOX['vt']} {data_fmt} {BOX['vt']}")
+
+
+def _make_bar(used: float, total: float, width: int) -> str:
+    if total <= 0:
+        return "░" * width
+    frac = min(max(used / total, 0.0), 1.0)
+    filled = int(round(frac * width))
+    filled = max(0, min(width, filled))
+    return "█" * filled + "░" * (width - filled)
+
+
+def _format_short_number(n: float) -> Optional[str]:
+    n = float(n)
+    for unit in ("", "K", "M", "B", "T"):
+        if abs(n) < 1000.0:
+            return f"{int(n)}" if unit == "" else f"{n:.1f}{unit}"
+        n /= 1000.0
+
+
+def _bytes_to_gb(nbytes: int) -> float:
+    return nbytes / (1024**3)
+
+
+def print_report_for_transformers(
+    model_id: str, revision: str, metadata: Dict[str, Any]
+) -> None:
+    ppdt = {}
+    for key, value in metadata.items():
+        if key in {"__metadata__"}:
+            continue
+        if value["dtype"] not in ppdt:
+            ppdt[value["dtype"]] = (0, 0)
+
+        match value["dtype"]:
+            case "F64" | "I64" | "U64":
+                dtype_b = 8
+            case "F32" | "I32" | "U32":
+                dtype_b = 4
+            case "F16" | "BF16" | "I16" | "U16":
+                dtype_b = 2
+            case "F8_E5M2" | "F8_E4M3" | "I8" | "U8":
+                dtype_b = 1
+            case _:
+                raise RuntimeError(f"DTYPE={value['dtype']} NOT HANDLED")
+
+        current_shape = math.prod(value["shape"])
+        current_shape_bytes = current_shape * dtype_b
+
+        ppdt[value["dtype"]] = (
+            ppdt[value["dtype"]][0] + current_shape,
+            ppdt[value["dtype"]][1] + current_shape_bytes,
+        )
+
+    rows = [
+        "INFERENCE MEMORY ESTIMATE FOR",
+        f"https://hf.co/{model_id} @ {revision}",
+        "TOTAL MEMORY",
+        "REQUIREMENTS",
+    ]
+    for dt, (params, nbytes) in ppdt.items():
+        rows.append(f"{dt} {params} {nbytes}")
+
+    max_len = 0
+    for r in rows:
+        max_len = max(max_len, len(str(r)))
+
+    current_len = min(max_len, MAX_DATA_LEN)
+
+    total_bytes = sum(nbytes for _, nbytes in ppdt.values())
+    total_params = sum(params for params, _ in ppdt.values())
+    total_gb = _bytes_to_gb(total_bytes)
+
+    _print_header(current_len)
+    _print_centered("INFERENCE MEMORY ESTIMATE FOR", current_len)
+    _print_centered(f"https://hf.co/{model_id} @ {revision}", current_len)
+    _print_divider(current_len + 1, "top")
+
+    total_text = f"{_bytes_to_gb(total_bytes):.2f} GB ({_format_short_number(total_params)} params)"
+    _print_row("TOTAL MEMORY", total_text, current_len)
+
+    total_bar = _make_bar(total_bytes, total_bytes, current_len)
+    _print_row("REQUIREMENTS", total_bar, current_len)
+    _print_divider(current_len + 1)
+
+    max_length = max([
+        len(f"{_format_short_number(params)} PARAMS") for params, _ in ppdt.values()
+    ])
+    for i, (dtype, (params, nbytes)) in enumerate(ppdt.items()):
+        dtype_name = dtype.upper()
+        dtype_gb = _bytes_to_gb(nbytes)
+
+        gb_text = f"{dtype_gb:.1f} / {total_gb:.1f} GB"
+        _print_row(
+            dtype_name + " " * (max_length - len(dtype_name)),
+            gb_text,
+            current_len,
+        )
+
+        bar = _make_bar(dtype_gb, total_gb, current_len)
+        _print_row(f"{_format_short_number(params)} PARAMS", bar, current_len)
+
+        if i < len(ppdt) - 1:
+            _print_divider(current_len + 1)
+
+    _print_divider(current_len + 1, "bottom")
+
+
+def print_report_for_diffusers(
+    model_id: str, revision: str, metadata: Dict[str, Dict[str, Any]]
+) -> None:
+    components_ppdt: Dict[str, Dict[str, tuple[int, int]]] = {}
+    total_bytes = 0
+    total_params = 0
+
+    for path, path_metadata in metadata.items():
+        ppdt: Dict[str, tuple[int, int]] = {}
+        for key, value in path_metadata.items():
+            if key in {"__metadata__"}:
+                continue
+
+            dtype = value["dtype"]
+            match dtype:
+                case "F64" | "I64" | "U64":
+                    dtype_b = 8
+                case "F32" | "I32" | "U32":
+                    dtype_b = 4
+                case "F16" | "BF16" | "I16" | "U16":
+                    dtype_b = 2
+                case "F8_E5M2" | "F8_E4M3" | "I8" | "U8":
+                    dtype_b = 1
+                case _:
+                    raise RuntimeError(f"DTYPE={dtype} NOT HANDLED")
+
+            current_shape = math.prod(value["shape"])
+            current_shape_bytes = current_shape * dtype_b
+
+            if dtype not in ppdt:
+                ppdt[dtype] = (0, 0)
+            ppdt[dtype] = (
+                ppdt[dtype][0] + current_shape,
+                ppdt[dtype][1] + current_shape_bytes,
+            )
+
+        components_ppdt[path] = ppdt
+        for params, nbytes in ppdt.values():
+            total_params += params
+            total_bytes += nbytes
+
+    rows = [
+        "INFERENCE MEMORY ESTIMATE FOR",
+        f"https://hf.co/{model_id} @ {revision}",
+        "TOTAL MEMORY",
+        "REQUIREMENTS",
+    ]
+    for path, ppdt in components_ppdt.items():
+        rows.append(path)
+        for dt, (params, nbytes) in ppdt.items():
+            rows.append(f"{dt} {params} {nbytes}")
+
+    max_len = 0
+    for r in rows:
+        max_len = max(max_len, len(str(r)))
+    current_len = min(max_len, MAX_DATA_LEN)
+
+    _print_header(current_len)
+    _print_centered("INFERENCE MEMORY ESTIMATE FOR", current_len)
+    _print_centered(f"https://hf.co/{model_id} @ {revision}", current_len)
+    _print_divider(current_len + 1, "top")
+
+    total_text = f"{_bytes_to_gb(total_bytes):.2f} GB ({_format_short_number(total_params)} params)"
+    _print_row("TOTAL MEMORY", total_text, current_len)
+
+    total_bar = _make_bar(total_bytes, total_bytes, current_len)
+    _print_row("REQUIREMENTS", total_bar, current_len)
+
+    for path, ppdt in components_ppdt.items():
+        _print_divider(current_len + 1, "top-continue")
+
+        path_bytes = sum(nbytes for _, nbytes in ppdt.values())
+        path_gb = _bytes_to_gb(path_bytes)
+
+        _print_centered(f"{path.upper()} ({path_gb:.2f} GB)", current_len)
+        _print_divider(current_len + 1, "top")
+
+        max_length = max([
+            len(f"{_format_short_number(params)} PARAMS") for params, _ in ppdt.values()
+        ])
+        for i, (dtype, (params, nbytes)) in enumerate(ppdt.items()):
+            dtype_name = dtype.upper()
+            dtype_gb = _bytes_to_gb(nbytes)
+
+            gb_text = f"{dtype_gb:.1f} / {_bytes_to_gb(total_bytes):.1f} GB"
+            _print_row(
+                dtype_name + " " * (max_length - len(dtype_name)),
+                gb_text,
+                current_len,
+            )
+
+            bar = _make_bar(dtype_gb, _bytes_to_gb(total_bytes), current_len)
+            _print_row(f"{_format_short_number(params)} PARAMS", bar, current_len)
+
+            if i < len(ppdt) - 1:
+                _print_divider(current_len + 1, "bottom-continue")
+
+    _print_divider(current_len + 1, "bottom")

hf_mem-0.0.3.dist-info/METADATA

@@ -0,0 +1,36 @@
+Metadata-Version: 2.4
+Name: hf-mem
+Version: 0.0.3
+Summary: A CLI to estimate inference memory requirements for Hugging Face models, written in Python
+Author: Alvaro Bartolome
+Author-email: Alvaro Bartolome <36760800+alvarobartt@users.noreply.github.com>
+License-Expression: MIT
+License-File: LICENSE
+Requires-Dist: httpx[http2]>=0.28.1
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+
+# `hf-mem`
+
+`hf-mem` is an experimental CLI to estimate inference memory requirements for Hugging Face models, written in Python. `hf-mem` is lightweight, only depends on `httpx`. It's recommended to run with [`uv`](https://github.com/astral-sh/uv) for a better experience.
+
+`hf-mem` lets you estimate the inference requirements to run any model from the Hugging Face Hub, including Transformers, Diffusers and Sentence Transformers models, as well as any model that contains [Safetensors](https://github.com/huggingface/safetensors) compatible weights.
+
+## Usage
+
+```bash
+uvx hf-mem --model-id MiniMaxAI/MiniMax-M2
+```
+
+<img width="899" height="648" src="https://github.com/user-attachments/assets/530f8b14-a415-4fd6-9054-bcd81cafae09" />
+
+```bash
+uvx hf-mem --model-id Qwen/Qwen-Image
+```
+
+<img width="894" height="872" src="https://github.com/user-attachments/assets/cd4234ec-bdcc-4db4-8b01-0ac9b5cd390c" />
+
+## References
+
+- [Safetensors Metadata parsing](https://huggingface.co/docs/safetensors/en/metadata_parsing)
+- [usgraphics - TR-100 Machine Report](https://github.com/usgraphics/usgc-machine-report)

hf_mem-0.0.3.dist-info/RECORD

@@ -0,0 +1,9 @@
+hf_mem/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+hf_mem/cli.py,sha256=m6W32NBWYvfqU_Odh8N0YAQHQ53P3HIhCdw90aICHAI,7781
+hf_mem/print.py,sha256=G3bDy2ND1WY8LyGTYf07E4zFZE54TJZu1uhVWGFnFkg,9269
+hf_mem/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+hf_mem-0.0.3.dist-info/licenses/LICENSE,sha256=TUO4cSlmwlYsHZfepsmo1zLoYhmT0RImXoRihVYNHhw,1060
+hf_mem-0.0.3.dist-info/WHEEL,sha256=M6du7VZflc4UPsGphmOXHANdgk8zessdJG0DBUuoA-U,78
+hf_mem-0.0.3.dist-info/entry_points.txt,sha256=p6o2afl5zELLNtkKwQxfLdOpal2nRlcshDZSfV02pOY,44
+hf_mem-0.0.3.dist-info/METADATA,sha256=n9wvVCPHGzSWC86iVJIH2jSS8hvovUAvjC-mB9o43Ac,1512
+hf_mem-0.0.3.dist-info/RECORD,,

hf_mem-0.0.3.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.9.5
+Root-Is-Purelib: true
+Tag: py3-none-any

hf_mem-0.0.3.dist-info/entry_points.txt

@@ -0,0 +1,3 @@
+[console_scripts]
+hf-mem = hf_mem.cli:main
+

hf_mem-0.0.3.dist-info/licenses/LICENSE

@@ -0,0 +1,25 @@
+Copyright (c) 2025 Alvaro Bartolome
+
+Permission is hereby granted, free of charge, to any
+person obtaining a copy of this software and associated
+documentation files (the "Software"), to deal in the
+Software without restriction, including without
+limitation the rights to use, copy, modify, merge,
+publish, distribute, sublicense, and/or sell copies of
+the Software, and to permit persons to whom the Software
+is furnished to do so, subject to the following
+conditions:
+
+The above copyright notice and this permission notice
+shall be included in all copies or substantial portions
+of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
+ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
+TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
+PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
+SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
+IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+DEALINGS IN THE SOFTWARE.