hf-mem 0.5.6__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
hf_mem-0.5.6/LICENSE ADDED
@@ -0,0 +1,25 @@
1
+ Copyright (c) 2026-present Alvaro Bartolome <alvarobartt@gmail.com>
2
+
3
+ Permission is hereby granted, free of charge, to any
4
+ person obtaining a copy of this software and associated
5
+ documentation files (the "Software"), to deal in the
6
+ Software without restriction, including without
7
+ limitation the rights to use, copy, modify, merge,
8
+ publish, distribute, sublicense, and/or sell copies of
9
+ the Software, and to permit persons to whom the Software
10
+ is furnished to do so, subject to the following
11
+ conditions:
12
+
13
+ The above copyright notice and this permission notice
14
+ shall be included in all copies or substantial portions
15
+ of the Software.
16
+
17
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF
18
+ ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
19
+ TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
20
+ PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
21
+ SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
22
+ CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
23
+ OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR
24
+ IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
25
+ DEALINGS IN THE SOFTWARE.
hf_mem-0.5.6/PKG-INFO ADDED
@@ -0,0 +1,131 @@
1
+ Metadata-Version: 2.4
2
+ Name: hf-mem
3
+ Version: 0.5.6
4
+ Summary: A CLI to estimate inference memory requirements for Hugging Face models, written in Python.
5
+ Author: Alvaro Bartolome
6
+ Author-email: Alvaro Bartolome <alvarobartt@gmail.com>
7
+ License-Expression: MIT
8
+ License-File: LICENSE
9
+ Requires-Dist: httpx2[http2]>=2.12.0
10
+ Requires-Python: >=3.10
11
+ Project-URL: Homepage, https://github.com/alvarobartt/hf-mem
12
+ Project-URL: Repository, https://github.com/alvarobartt/hf-mem
13
+ Project-URL: Issues, https://github.com/alvarobartt/hf-mem/issues
14
+ Description-Content-Type: text/markdown
15
+
16
+ <img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/hf-mem.png" />
17
+
18
+ ---
19
+
20
+ > [!WARNING]
21
+ > `hf-mem` is still experimental and therefore subject to major changes across releases, so please keep in mind that breaking changes may occur until v1.0.0.
22
+
23
+ `hf-mem` is a CLI to estimate inference memory requirements for Hugging Face models, written in Python. `hf-mem` is lightweight, only depends on [`httpx2`](https://github.com/pydantic/httpx2), as it pulls the [Safetensors](https://github.com/huggingface/safetensors) and / or [GGUF](https://github.com/ggml-org/ggml/blob/master/docs/gguf.md) metadata via [HTTP Range requests](https://developer.mozilla.org/en-US/docs/Web/HTTP/Guides/Range_requests). It's recommended to run with [`uv`](https://github.com/astral-sh/uv) for a better experience.
24
+
25
+ `hf-mem` lets you estimate the inference requirements to run any model from the Hugging Face Hub, including [Transformers](https://github.com/huggingface/transformers), [Diffusers](https://github.com/huggingface/diffusers) and [Sentence Transformers](https://github.com/huggingface/sentence-transformers) models, or really any model as long as it contains any of [Safetensors](https://github.com/huggingface/safetensors) or [GGUF](https://github.com/ggml-org/ggml/blob/master/docs/gguf.md) weights.
26
+
27
+ Read more information about `hf-mem` in [this short-form post](https://alvarobartt.com/hf-mem), but note it's not up-to-date as it was written in January 2026.
28
+
29
+ ## Usage
30
+
31
+ ### CLI (Recommended)
32
+
33
+ #### Transformers
34
+
35
+ ```bash
36
+ uvx hf-mem --model-id MiniMaxAI/MiniMax-M2
37
+ ```
38
+
39
+ <img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/transformers.png" />
40
+
41
+ #### Diffusers
42
+
43
+ ```bash
44
+ uvx hf-mem --model-id Qwen/Qwen-Image
45
+ ```
46
+
47
+ <img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/diffusers.png" />
48
+
49
+ #### Sentence Transformers
50
+
51
+ ```bash
52
+ uvx hf-mem --model-id google/embeddinggemma-300m
53
+ ```
54
+
55
+ <img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/sentence-transformers.png" />
56
+
57
+ ### Python
58
+
59
+ You can also run it programmatically with Python as:
60
+
61
+ ```python
62
+ from hf_mem import run
63
+
64
+ result = run(model_id="MiniMaxAI/MiniMax-M2", experimental=True)
65
+ print(result)
66
+ # Result(model_id='MiniMaxAI/MiniMax-M2', revision='main', filename=None, memory=230121630720, kv_cache=24964497408, total_memory=255086128128, details=False)
67
+ ```
68
+
69
+ If you're already inside an async application, use `arun(...)` instead:
70
+
71
+ ```python
72
+ from hf_mem import arun
73
+
74
+ result = await arun(model_id="MiniMaxAI/MiniMax-M2", experimental=True)
75
+ print(result)
76
+ # Result(model_id='MiniMaxAI/MiniMax-M2', revision='main', filename=None, memory=230121630720, kv_cache=24964497408, total_memory=255086128128, details=False)
77
+ ```
78
+
79
+ ## Experimental
80
+
81
+ By enabling the `--experimental` flag, you can enable the KV Cache memory estimation for LLMs (`...ForCausalLM`) and VLMs (`...ForConditionalGeneration`), even including a custom `--max-model-len` (defaults to the `config.json` default), `--batch-size` (defaults to 1), and the `--kv-cache-dtype` (defaults to `auto` which means it uses the default data type set in `config.json` under `torch_dtype` or `dtype`, or rather from `quantization_config` when applicable).
82
+
83
+ Additionally, for Mixture of Experts (MoEs), the `--experimental` flag also adds a weights-only breakdown of the base model and all the experts combined.
84
+
85
+ ```bash
86
+ uvx hf-mem --model-id MiniMaxAI/MiniMax-M2 --experimental
87
+ ```
88
+
89
+ <img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/experimental.png" />
90
+
91
+ ## GGUF
92
+
93
+ If the repository contains GGUF model weights, those will be listed by default (only if there are no Safetensors weights, otherwise the GGUFs will be ignored) and the memory will be estimated for each one of those; whereas if a specific file is provided, then the memory estimation will be targeted for that given file instead.
94
+
95
+ ```bash
96
+ uvx hf-mem --model-id TheBloke/deepseek-llm-7B-chat-GGUF --experimental
97
+ ```
98
+
99
+ <img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/gguf.png" />
100
+
101
+ Or if you want to only get the estimation on a given file:
102
+
103
+ ```bash
104
+ uvx hf-mem --model-id TheBloke/deepseek-llm-7B-chat-GGUF --gguf-file deepseek-llm-7b-chat.Q2_K.gguf --experimental
105
+ ```
106
+
107
+ <img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/gguf-file.png" />
108
+
109
+ ## Skills
110
+
111
+ Optionally, you can add `hf-mem` as an agent skill, which allows the underlying coding agent to discover and use it when provided as a [`SKILL.md`](skills/hf-mem/SKILL.md), e.g., `.claude/skills/hf-mem/SKILL.md`.
112
+
113
+ More information can be found at [Anthropic Agent Skills and how to use them](https://github.com/anthropics/skills).
114
+
115
+ ## Extensions
116
+
117
+ Optionally, you can also add `hf-mem` as an extension to the Hugging Face Hub CLI under `hf` as `hf mem ...`. First you need to install `hf` as explained in [Hugging Face CLI - "Getting started"](https://huggingface.co/docs/huggingface_hub/en/guides/cli). To add `hf-mem` as an extension, all you need to do is run:
118
+
119
+ ```bash
120
+ hf extensions add alvarobartt/hf-mem
121
+ ```
122
+ More information can be found at [Hugging Face Hub CLI - "How Python extensions are installed"](https://huggingface.co/docs/huggingface_hub/en/guides/cli-extensions#how-python-extensions-are-installed).
123
+
124
+ Once installed, you can use `hf-mem` via `hf` as `hf mem ...` along with the rest of Hugging Face extensions you have installed (which you can list via `hf extensions list`).
125
+
126
+ ## References
127
+
128
+ - [Safetensors Metadata parsing](https://huggingface.co/docs/safetensors/en/metadata_parsing)
129
+ - [GGUF File format](https://github.com/ggml-org/ggml/blob/master/docs/gguf.md)
130
+ - [GGUF on the Hugging Face Hub](https://huggingface.co/docs/hub/en/gguf)
131
+ - [usgraphics - TR-100 Machine Report](https://github.com/usgraphics/usgc-machine-report)
hf_mem-0.5.6/README.md ADDED
@@ -0,0 +1,116 @@
1
+ <img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/hf-mem.png" />
2
+
3
+ ---
4
+
5
+ > [!WARNING]
6
+ > `hf-mem` is still experimental and therefore subject to major changes across releases, so please keep in mind that breaking changes may occur until v1.0.0.
7
+
8
+ `hf-mem` is a CLI to estimate inference memory requirements for Hugging Face models, written in Python. `hf-mem` is lightweight, only depends on [`httpx2`](https://github.com/pydantic/httpx2), as it pulls the [Safetensors](https://github.com/huggingface/safetensors) and / or [GGUF](https://github.com/ggml-org/ggml/blob/master/docs/gguf.md) metadata via [HTTP Range requests](https://developer.mozilla.org/en-US/docs/Web/HTTP/Guides/Range_requests). It's recommended to run with [`uv`](https://github.com/astral-sh/uv) for a better experience.
9
+
10
+ `hf-mem` lets you estimate the inference requirements to run any model from the Hugging Face Hub, including [Transformers](https://github.com/huggingface/transformers), [Diffusers](https://github.com/huggingface/diffusers) and [Sentence Transformers](https://github.com/huggingface/sentence-transformers) models, or really any model as long as it contains any of [Safetensors](https://github.com/huggingface/safetensors) or [GGUF](https://github.com/ggml-org/ggml/blob/master/docs/gguf.md) weights.
11
+
12
+ Read more information about `hf-mem` in [this short-form post](https://alvarobartt.com/hf-mem), but note it's not up-to-date as it was written in January 2026.
13
+
14
+ ## Usage
15
+
16
+ ### CLI (Recommended)
17
+
18
+ #### Transformers
19
+
20
+ ```bash
21
+ uvx hf-mem --model-id MiniMaxAI/MiniMax-M2
22
+ ```
23
+
24
+ <img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/transformers.png" />
25
+
26
+ #### Diffusers
27
+
28
+ ```bash
29
+ uvx hf-mem --model-id Qwen/Qwen-Image
30
+ ```
31
+
32
+ <img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/diffusers.png" />
33
+
34
+ #### Sentence Transformers
35
+
36
+ ```bash
37
+ uvx hf-mem --model-id google/embeddinggemma-300m
38
+ ```
39
+
40
+ <img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/sentence-transformers.png" />
41
+
42
+ ### Python
43
+
44
+ You can also run it programmatically with Python as:
45
+
46
+ ```python
47
+ from hf_mem import run
48
+
49
+ result = run(model_id="MiniMaxAI/MiniMax-M2", experimental=True)
50
+ print(result)
51
+ # Result(model_id='MiniMaxAI/MiniMax-M2', revision='main', filename=None, memory=230121630720, kv_cache=24964497408, total_memory=255086128128, details=False)
52
+ ```
53
+
54
+ If you're already inside an async application, use `arun(...)` instead:
55
+
56
+ ```python
57
+ from hf_mem import arun
58
+
59
+ result = await arun(model_id="MiniMaxAI/MiniMax-M2", experimental=True)
60
+ print(result)
61
+ # Result(model_id='MiniMaxAI/MiniMax-M2', revision='main', filename=None, memory=230121630720, kv_cache=24964497408, total_memory=255086128128, details=False)
62
+ ```
63
+
64
+ ## Experimental
65
+
66
+ By enabling the `--experimental` flag, you can enable the KV Cache memory estimation for LLMs (`...ForCausalLM`) and VLMs (`...ForConditionalGeneration`), even including a custom `--max-model-len` (defaults to the `config.json` default), `--batch-size` (defaults to 1), and the `--kv-cache-dtype` (defaults to `auto` which means it uses the default data type set in `config.json` under `torch_dtype` or `dtype`, or rather from `quantization_config` when applicable).
67
+
68
+ Additionally, for Mixture of Experts (MoEs), the `--experimental` flag also adds a weights-only breakdown of the base model and all the experts combined.
69
+
70
+ ```bash
71
+ uvx hf-mem --model-id MiniMaxAI/MiniMax-M2 --experimental
72
+ ```
73
+
74
+ <img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/experimental.png" />
75
+
76
+ ## GGUF
77
+
78
+ If the repository contains GGUF model weights, those will be listed by default (only if there are no Safetensors weights, otherwise the GGUFs will be ignored) and the memory will be estimated for each one of those; whereas if a specific file is provided, then the memory estimation will be targeted for that given file instead.
79
+
80
+ ```bash
81
+ uvx hf-mem --model-id TheBloke/deepseek-llm-7B-chat-GGUF --experimental
82
+ ```
83
+
84
+ <img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/gguf.png" />
85
+
86
+ Or if you want to only get the estimation on a given file:
87
+
88
+ ```bash
89
+ uvx hf-mem --model-id TheBloke/deepseek-llm-7B-chat-GGUF --gguf-file deepseek-llm-7b-chat.Q2_K.gguf --experimental
90
+ ```
91
+
92
+ <img src="https://huggingface.co/datasets/alvarobartt/hf-mem/resolve/main/gguf-file.png" />
93
+
94
+ ## Skills
95
+
96
+ Optionally, you can add `hf-mem` as an agent skill, which allows the underlying coding agent to discover and use it when provided as a [`SKILL.md`](skills/hf-mem/SKILL.md), e.g., `.claude/skills/hf-mem/SKILL.md`.
97
+
98
+ More information can be found at [Anthropic Agent Skills and how to use them](https://github.com/anthropics/skills).
99
+
100
+ ## Extensions
101
+
102
+ Optionally, you can also add `hf-mem` as an extension to the Hugging Face Hub CLI under `hf` as `hf mem ...`. First you need to install `hf` as explained in [Hugging Face CLI - "Getting started"](https://huggingface.co/docs/huggingface_hub/en/guides/cli). To add `hf-mem` as an extension, all you need to do is run:
103
+
104
+ ```bash
105
+ hf extensions add alvarobartt/hf-mem
106
+ ```
107
+ More information can be found at [Hugging Face Hub CLI - "How Python extensions are installed"](https://huggingface.co/docs/huggingface_hub/en/guides/cli-extensions#how-python-extensions-are-installed).
108
+
109
+ Once installed, you can use `hf-mem` via `hf` as `hf mem ...` along with the rest of Hugging Face extensions you have installed (which you can list via `hf extensions list`).
110
+
111
+ ## References
112
+
113
+ - [Safetensors Metadata parsing](https://huggingface.co/docs/safetensors/en/metadata_parsing)
114
+ - [GGUF File format](https://github.com/ggml-org/ggml/blob/master/docs/gguf.md)
115
+ - [GGUF on the Hugging Face Hub](https://huggingface.co/docs/hub/en/gguf)
116
+ - [usgraphics - TR-100 Machine Report](https://github.com/usgraphics/usgc-machine-report)
@@ -0,0 +1,40 @@
1
+ [project]
2
+ name = "hf-mem"
3
+ version = "0.5.6"
4
+ description = "A CLI to estimate inference memory requirements for Hugging Face models, written in Python."
5
+ readme = "README.md"
6
+ license = "MIT"
7
+ license-files = ["LICENSE"]
8
+ requires-python = ">=3.10"
9
+ dependencies = ["httpx2[http2]>=2.12.0"]
10
+
11
+ [[project.authors]]
12
+ name = "Alvaro Bartolome"
13
+ email = "alvarobartt@gmail.com"
14
+
15
+ [project.urls]
16
+ Homepage = "https://github.com/alvarobartt/hf-mem"
17
+ Repository = "https://github.com/alvarobartt/hf-mem"
18
+ Issues = "https://github.com/alvarobartt/hf-mem/issues"
19
+
20
+ [project.scripts]
21
+ hf-mem = "hf_mem.cli:main"
22
+
23
+ [build-system]
24
+ requires = ["uv_build>=0.9.5,<0.10.0"]
25
+ build-backend = "uv_build"
26
+
27
+ [dependency-groups]
28
+ lint = [
29
+ "pre-commit>=4.5.1",
30
+ "ruff>=0.14.10",
31
+ ]
32
+
33
+ [tool.ruff]
34
+ line-length = 112
35
+
36
+ [tool.ruff.lint]
37
+ select = ["I"]
38
+
39
+ [tool.ruff.lint.isort]
40
+ force-single-line = false
@@ -0,0 +1,41 @@
1
+ [project]
2
+ name = "hf-mem"
3
+ version = "0.5.6"
4
+ description = "A CLI to estimate inference memory requirements for Hugging Face models, written in Python."
5
+ readme = "README.md"
6
+ authors = [
7
+ { name = "Alvaro Bartolome", email = "alvarobartt@gmail.com" }
8
+ ]
9
+ license = "MIT"
10
+ license-files = ["LICENSE"]
11
+ requires-python = ">=3.10"
12
+ dependencies = [
13
+ "httpx2[http2]>=2.12.0",
14
+ ]
15
+
16
+ [project.urls]
17
+ Homepage = "https://github.com/alvarobartt/hf-mem"
18
+ Repository = "https://github.com/alvarobartt/hf-mem"
19
+ Issues = "https://github.com/alvarobartt/hf-mem/issues"
20
+
21
+ [build-system]
22
+ requires = ["uv_build>=0.9.5,<0.10.0"]
23
+ build-backend = "uv_build"
24
+
25
+ [dependency-groups]
26
+ lint = [
27
+ "pre-commit>=4.5.1",
28
+ "ruff>=0.14.10",
29
+ ]
30
+
31
+ [tool.ruff]
32
+ line-length = 112
33
+
34
+ [tool.ruff.lint]
35
+ select = ["I"]
36
+
37
+ [tool.ruff.lint.isort]
38
+ force-single-line = false
39
+
40
+ [project.scripts]
41
+ hf-mem = "hf_mem.cli:main"
@@ -0,0 +1,4 @@
1
+ from hf_mem._version import __version__
2
+ from hf_mem.run import KvCache, Result, arun, run
3
+
4
+ __all__ = ["__version__", "run", "arun", "Result", "KvCache"]
@@ -0,0 +1,14 @@
1
+ import os
2
+ from typing import Any, Dict
3
+
4
+ import httpx2
5
+
6
+ # NOTE: Shared across all HTTP fetch operations; overridable via the `REQUEST_TIMEOUT` environment variable
7
+ REQUEST_TIMEOUT = float(os.getenv("REQUEST_TIMEOUT", 30.0))
8
+
9
+
10
+ # NOTE: Return type-hint set to `Any`, but it will only be a JSON-compatible object
11
+ async def get_json_file(client: httpx2.AsyncClient, url: str, headers: Dict[str, str] | None = None) -> Any:
12
+ response = await client.get(url, headers=headers, timeout=REQUEST_TIMEOUT)
13
+ response.raise_for_status()
14
+ return response.json()
@@ -0,0 +1,133 @@
1
+ from typing import Literal
2
+
3
+ MIN_NAME_LEN = 5
4
+ MAX_NAME_LEN = 14
5
+ MAX_DATA_LEN = 64
6
+ BORDERS_AND_PADDING = 4
7
+
8
+ BOX = {
9
+ "tl": "┌",
10
+ "tr": "┐",
11
+ "bl": "└",
12
+ "br": "┘",
13
+ "ht": "─",
14
+ "vt": "│",
15
+ "tsep": "┬",
16
+ "bsep": "┴",
17
+ "lm": "├",
18
+ "rm": "┤",
19
+ "mm": "┼",
20
+ }
21
+
22
+
23
+ def _print_with_color(content: str) -> None:
24
+ print(f"\x1b[38;2;244;183;63m{content}\x1b[0m")
25
+
26
+
27
+ def _print_header(current_len: int, badge: str | None = None) -> None:
28
+ length = current_len + 2 * BORDERS_AND_PADDING + 2
29
+ if badge:
30
+ badge_text = f" {badge} "
31
+ fill_len = max(0, length - 2 - len(badge_text))
32
+ top = BOX["tl"] + (BOX["tsep"] * fill_len) + badge_text + BOX["tr"]
33
+ else:
34
+ top = BOX["tl"] + (BOX["tsep"] * (length - 2)) + BOX["tr"]
35
+ _print_with_color(top)
36
+
37
+ bottom = BOX["lm"] + (BOX["bsep"] * (length - 2)) + BOX["rm"]
38
+ _print_with_color(bottom)
39
+
40
+
41
+ def _print_centered(text: str, current_len: int) -> None:
42
+ total_width = current_len + 2 * BORDERS_AND_PADDING
43
+ text_len = len(text)
44
+ pad_left = (total_width - text_len) // 2
45
+ pad_right = total_width - text_len - pad_left
46
+ _print_with_color(f"{BOX['vt']}{' ' * pad_left}{text}{' ' * pad_right}{BOX['vt']}")
47
+
48
+
49
+ def _print_divider(
50
+ current_len: int,
51
+ side: Literal["top", "top-continue", "mid", "bottom", "bottom-continue"] | None = None,
52
+ name_len: int = MAX_NAME_LEN,
53
+ ) -> None:
54
+ match side:
55
+ case "top":
56
+ left, mid, right = BOX["lm"], BOX["tsep"], BOX["rm"]
57
+ case "top-continue":
58
+ left, mid, right = BOX["lm"], BOX["bsep"], BOX["rm"]
59
+ case "bottom":
60
+ left, mid, right = BOX["bl"], BOX["bsep"], BOX["br"]
61
+ case "bottom-continue":
62
+ left, mid, right = BOX["lm"], BOX["bsep"], BOX["rm"]
63
+ case "mid":
64
+ left, mid, right = BOX["lm"], BOX["ht"], BOX["rm"]
65
+ case _:
66
+ left, mid, right = BOX["lm"], BOX["mm"], BOX["rm"]
67
+
68
+ name_col_inner = name_len + 2
69
+ data_col_inner = current_len + 1
70
+
71
+ line = left
72
+ line += BOX["ht"] * name_col_inner
73
+ line += mid
74
+ line += BOX["ht"] * data_col_inner
75
+ line += right
76
+ _print_with_color(line)
77
+
78
+
79
+ def _print_full_divider(
80
+ current_len: int,
81
+ side: Literal["top", "top-continue", "bottom", "bottom-continue", "mid"] | None = None,
82
+ ) -> None:
83
+ total_inner = current_len + 2 * BORDERS_AND_PADDING
84
+
85
+ match side:
86
+ case "top":
87
+ left, right = BOX["lm"], BOX["rm"]
88
+ case "top-continue":
89
+ left, right = BOX["lm"], BOX["rm"]
90
+ case "bottom":
91
+ left, right = BOX["bl"], BOX["br"]
92
+ case "bottom-continue":
93
+ left, right = BOX["lm"], BOX["rm"]
94
+ case _:
95
+ left, right = BOX["lm"], BOX["rm"]
96
+
97
+ _print_with_color(left + BOX["ht"] * total_inner + right)
98
+
99
+
100
+ def _format_name(name: str, max_len: int = MAX_NAME_LEN) -> str:
101
+ if len(name) < MIN_NAME_LEN:
102
+ return f"{name:<{MIN_NAME_LEN}}"
103
+ if len(name) > max_len:
104
+ return name[: max_len - 3] + "..."
105
+ return f"{name:<{max_len}}"
106
+
107
+
108
+ def _print_row(name: str, text: str, current_len: int, name_len: int = MAX_NAME_LEN) -> None:
109
+ name_fmt = _format_name(name, name_len)
110
+ data_fmt = f"{str(text):<{current_len}}"
111
+ _print_with_color(f"{BOX['vt']} {name_fmt} {BOX['vt']} {data_fmt} {BOX['vt']}")
112
+
113
+
114
+ def _make_bar(used: float, total: float, width: int) -> str:
115
+ if total <= 0:
116
+ return "░" * width
117
+ frac = min(max(used / total, 0.0), 1.0)
118
+ filled = int(round(frac * width))
119
+ filled = max(0, min(width, filled))
120
+ return "█" * filled + "░" * (width - filled)
121
+
122
+
123
+ def _format_short_number(n: float) -> str:
124
+ n = float(n)
125
+ for unit in ("", "K", "M", "B", "T"):
126
+ if abs(n) < 1000.0:
127
+ return f"{int(n)}" if unit == "" else f"{n:.2f}{unit}"
128
+ n /= 1000.0
129
+ return f"{n:.2f}P"
130
+
131
+
132
+ def _bytes_to_gib(nbytes: int) -> float:
133
+ return nbytes / (1024**3)
@@ -0,0 +1,9 @@
1
+ from dataclasses import dataclass
2
+
3
+
4
+ @dataclass
5
+ class KvCache:
6
+ max_model_len: int
7
+ cache_size: int
8
+ batch_size: int
9
+ cache_dtype: str | None = None
@@ -0,0 +1,5 @@
1
+ from importlib.metadata import version
2
+
3
+ # NOTE: Single source of truth for the package version, imported by all sub-modules that need it
4
+ # to avoid circular imports through hf_mem.__init__
5
+ __version__: str = version("hf-mem")
@@ -0,0 +1,148 @@
1
+ import argparse
2
+ import asyncio
3
+ import json
4
+ import warnings
5
+
6
+ from hf_mem.gguf.print import print_gguf_files_report, print_gguf_report
7
+ from hf_mem.gguf.types import GGUFDtype
8
+ from hf_mem.run import Result, arun
9
+ from hf_mem.safetensors.kv_cache import KV_CACHE_DTYPE_CHOICES
10
+ from hf_mem.safetensors.print import print_safetensors_report
11
+
12
+ VALID_KV_CACHE_DTYPES = set(KV_CACHE_DTYPE_CHOICES) | set(GGUFDtype.__members__.keys())
13
+
14
+
15
+ def _print_result(result: Result) -> None:
16
+ if result.safetensors is not None:
17
+ print_safetensors_report(
18
+ model_id=result.model_id,
19
+ revision=result.revision,
20
+ metadata=result.safetensors,
21
+ kv_cache=result.kv_cache_metadata,
22
+ moe=result.moe_metadata,
23
+ )
24
+ return
25
+
26
+ if result.gguf_files is not None:
27
+ if result.filename is not None:
28
+ print_gguf_report(
29
+ model_id=result.model_id,
30
+ filename=result.filename,
31
+ revision=result.revision,
32
+ metadata=result.gguf_files[result.filename],
33
+ kv_cache=result.gguf_files[result.filename].kv_cache,
34
+ )
35
+ else:
36
+ print_gguf_files_report(
37
+ model_id=result.model_id,
38
+ revision=result.revision,
39
+ gguf_files=result.gguf_files,
40
+ memory=result.memory, # type: ignore[arg-type]
41
+ kv_cache=result.kv_cache, # type: ignore[arg-type]
42
+ )
43
+
44
+
45
+ def main() -> None:
46
+ parser = argparse.ArgumentParser()
47
+
48
+ parser.add_argument("--model-id", required=True, help="Model ID on the Hugging Face Hub")
49
+ parser.add_argument(
50
+ "--revision",
51
+ default="main",
52
+ help="Model revision on the Hugging Face Hub",
53
+ )
54
+ parser.add_argument(
55
+ "--hf-token",
56
+ type=str,
57
+ default=None,
58
+ required=False,
59
+ help="The Hugging Face Hub token to use when sending requests to the Hub. If not provided it will be retrieved from either the `HF_TOKEN` environment variable or from the `.cache/huggingface/token` file if applicable.",
60
+ )
61
+
62
+ parser.add_argument(
63
+ "--experimental",
64
+ action="store_true",
65
+ help=(
66
+ "Whether to enable the experimental KV Cache estimation or not. For Safetensors "
67
+ "MoE Transformers models it also includes a base-model-vs-experts weights "
68
+ "breakdown. Only applies to `...ForCausalLM` and `...ForConditionalGeneration` "
69
+ "models from Transformers."
70
+ ),
71
+ )
72
+ parser.add_argument(
73
+ "--max-model-len",
74
+ type=int,
75
+ default=None,
76
+ # NOTE: https://docs.vllm.ai/en/stable/configuration/engine_args/#-max-model-len
77
+ help="Model context length (prompt and output). If unspecified, will be automatically derived from the model config.",
78
+ )
79
+ parser.add_argument(
80
+ "--batch-size",
81
+ type=int,
82
+ default=1,
83
+ help="Batch size to help estimate the required RAM for caching when running the inference. Defaults to 1.",
84
+ )
85
+ parser.add_argument(
86
+ "--kv-cache-dtype",
87
+ type=str,
88
+ default="auto",
89
+ # NOTE: https://docs.vllm.ai/en/stable/cli/serve/#-kv-cache-dtype
90
+ help=f"Data type for the KV cache storage. If `auto` is specified, it will use the default model dtype specified in the `config.json` (if available) or F16 for GGUF files. Despite the FP8 data types having different formats, all those take 1 byte, meaning that the calculation would lead to the same results. Valid values are {KV_CACHE_DTYPE_CHOICES} for safetensors files and {['auto'] + list(GGUFDtype.__members__.keys())} for GGUF files. Defaults to `auto`.",
91
+ )
92
+
93
+ parser.add_argument(
94
+ "--json-output",
95
+ action="store_true",
96
+ help="Whether to provide the output as a JSON instead of printed as table.",
97
+ )
98
+ parser.add_argument(
99
+ "--ignore-table-width",
100
+ action="store_true",
101
+ help="Deprecated: table width now auto-expands when needed, so this flag is ignored.",
102
+ )
103
+ parser.add_argument(
104
+ "--gguf-file",
105
+ type=str,
106
+ default=None,
107
+ help="Specific GGUF file to estimate. If not provided, all GGUF files found in the repo will be estimated. The path within the repo is required (e.g. Q4_K_M/model-Q4_K_M.gguf).",
108
+ )
109
+ parser.add_argument(
110
+ "--details",
111
+ action="store_true",
112
+ help="Include per-component dtype and parameter breakdowns in JSON output (`--json-output` only).",
113
+ )
114
+ args = parser.parse_args()
115
+
116
+ if args.experimental:
117
+ warnings.warn(
118
+ "`--experimental` is set, which means that models with an architecture as `...ForCausalLM` and `...ForConditionalGeneration` will include estimations for the KV Cache as well. You can also provide the args `--max-model-len` and `--batch-size` as part of the estimation. Note that enabling `--experimental` means that the output will be different both when displayed and when dumped as JSON with `--json-output`, so bear that in mind."
119
+ )
120
+
121
+ if args.ignore_table_width:
122
+ warnings.warn(
123
+ "`--ignore-table-width` is deprecated and has no effect; table width auto-expands when needed by default."
124
+ )
125
+
126
+ if args.kv_cache_dtype not in VALID_KV_CACHE_DTYPES:
127
+ raise RuntimeError(
128
+ f"--kv-cache-dtype={args.kv_cache_dtype} not recognized. Valid options: {VALID_KV_CACHE_DTYPES}."
129
+ )
130
+
131
+ result = asyncio.run(
132
+ arun(
133
+ model_id=args.model_id,
134
+ revision=args.revision,
135
+ hf_token=args.hf_token,
136
+ experimental=args.experimental,
137
+ max_model_len=args.max_model_len,
138
+ batch_size=args.batch_size,
139
+ kv_cache_dtype=args.kv_cache_dtype,
140
+ gguf_file=args.gguf_file,
141
+ details=args.details,
142
+ )
143
+ )
144
+
145
+ if args.json_output:
146
+ print(json.dumps(result.to_json()))
147
+ else:
148
+ _print_result(result)