hf-mem 0.2.2__py3-none-any.whl
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/__init__.py +0 -0
- hf_mem/cli.py +231 -0
- hf_mem/metadata.py +61 -0
- hf_mem/print.py +196 -0
- hf_mem/py.typed +0 -0
- hf_mem/types.py +32 -0
- hf_mem-0.2.2.dist-info/METADATA +40 -0
- hf_mem-0.2.2.dist-info/RECORD +11 -0
- hf_mem-0.2.2.dist-info/WHEEL +4 -0
- hf_mem-0.2.2.dist-info/entry_points.txt +3 -0
- hf_mem-0.2.2.dist-info/licenses/LICENSE +25 -0
hf_mem/__init__.py
ADDED
|
File without changes
|
hf_mem/cli.py
ADDED
|
@@ -0,0 +1,231 @@
|
|
|
1
|
+
import argparse
|
|
2
|
+
import asyncio
|
|
3
|
+
import json
|
|
4
|
+
import os
|
|
5
|
+
import struct
|
|
6
|
+
from dataclasses import asdict
|
|
7
|
+
from functools import reduce
|
|
8
|
+
from typing import Any, Dict, List, Optional
|
|
9
|
+
|
|
10
|
+
import httpx
|
|
11
|
+
|
|
12
|
+
from hf_mem.metadata import parse_safetensors_metadata
|
|
13
|
+
from hf_mem.print import print_report
|
|
14
|
+
|
|
15
|
+
# NOTE: Defines the bytes that will be fetched per safetensors file, but the metadata
|
|
16
|
+
# can indeed be larger than that
|
|
17
|
+
MAX_METADATA_SIZE = 100_000
|
|
18
|
+
REQUEST_TIMEOUT = float(os.getenv("REQUEST_TIMEOUT", 10.0))
|
|
19
|
+
MAX_CONCURRENCY = int(os.getenv("MAX_WORKERS", min(32, (os.cpu_count() or 1) + 4)))
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
# NOTE: Return type-hint set to `Any`, but it will only be a JSON-compatible object
|
|
23
|
+
async def get_json_file(client: httpx.AsyncClient, url: str, headers: Optional[Dict[str, str]] = None) -> Any:
|
|
24
|
+
response = await client.get(url, headers=headers, timeout=REQUEST_TIMEOUT)
|
|
25
|
+
response.raise_for_status()
|
|
26
|
+
return response.json()
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
async def fetch_safetensors_metadata(
|
|
30
|
+
client: httpx.AsyncClient, url: str, headers: Optional[Dict[str, str]] = None
|
|
31
|
+
) -> Dict[str, Any]:
|
|
32
|
+
headers = {"Range": f"bytes=0-{MAX_METADATA_SIZE}", **(headers or {})}
|
|
33
|
+
response = await client.get(url, headers=headers, timeout=REQUEST_TIMEOUT)
|
|
34
|
+
response.raise_for_status()
|
|
35
|
+
|
|
36
|
+
metadata = response.read()
|
|
37
|
+
# NOTE: Parse the first 8 bytes as a little-endian uint64 (size of the metadata)
|
|
38
|
+
metadata_size = struct.unpack("<Q", metadata[:8])[0]
|
|
39
|
+
|
|
40
|
+
if metadata_size < MAX_METADATA_SIZE:
|
|
41
|
+
metadata = metadata[8 : metadata_size + 8]
|
|
42
|
+
return json.loads(metadata)
|
|
43
|
+
|
|
44
|
+
# NOTE: Given that by default we just fetch the first 100_000 bytes, if the content is larger
|
|
45
|
+
# then we simply fetch the remainder again
|
|
46
|
+
metadata = metadata[8 : MAX_METADATA_SIZE + 8]
|
|
47
|
+
headers["Range"] = f"bytes={MAX_METADATA_SIZE + 1}-{metadata_size + 7}"
|
|
48
|
+
|
|
49
|
+
response = await client.get(url, headers=headers, timeout=REQUEST_TIMEOUT)
|
|
50
|
+
response.raise_for_status()
|
|
51
|
+
|
|
52
|
+
metadata += response.read()
|
|
53
|
+
return json.loads(metadata)
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
async def run(
|
|
57
|
+
model_id: str,
|
|
58
|
+
revision: str,
|
|
59
|
+
json_output: bool = False,
|
|
60
|
+
ignore_table_width: bool = False,
|
|
61
|
+
) -> Dict[str, Any] | None:
|
|
62
|
+
headers = {}
|
|
63
|
+
# NOTE: Read from `HF_TOKEN` if provided, then fallback to reading from `$HF_HOME/token`
|
|
64
|
+
if token := os.getenv("HF_TOKEN"):
|
|
65
|
+
headers["Authorization"] = f"Bearer {token}"
|
|
66
|
+
elif "Authorization" not in headers:
|
|
67
|
+
path = os.getenv("HF_HOME", ".cache/huggingface")
|
|
68
|
+
filename = (
|
|
69
|
+
os.path.join(os.path.expanduser("~"), path, "token")
|
|
70
|
+
if not os.path.isabs(path)
|
|
71
|
+
else os.path.join(path, "token")
|
|
72
|
+
)
|
|
73
|
+
|
|
74
|
+
if os.path.exists(filename):
|
|
75
|
+
with open(filename, "r", encoding="utf-8") as f:
|
|
76
|
+
headers["Authorization"] = f"Bearer {f.read().strip()}"
|
|
77
|
+
|
|
78
|
+
client = httpx.AsyncClient(
|
|
79
|
+
limits=httpx.Limits(
|
|
80
|
+
max_keepalive_connections=MAX_CONCURRENCY,
|
|
81
|
+
max_connections=MAX_CONCURRENCY,
|
|
82
|
+
),
|
|
83
|
+
timeout=httpx.Timeout(REQUEST_TIMEOUT),
|
|
84
|
+
http2=True,
|
|
85
|
+
follow_redirects=True,
|
|
86
|
+
)
|
|
87
|
+
|
|
88
|
+
# TODO: `recursive=true` shouldn't really be required unless it's a Diffusers
|
|
89
|
+
# models... I don't think this adds extra latency anyway
|
|
90
|
+
url = f"https://huggingface.co/api/models/{model_id}/tree/{revision}?recursive=true"
|
|
91
|
+
files = await get_json_file(client=client, url=url, headers=headers)
|
|
92
|
+
file_paths = [f["path"] for f in files if f.get("path") and f.get("type") == "file"]
|
|
93
|
+
|
|
94
|
+
if "model.safetensors" in file_paths:
|
|
95
|
+
url = f"https://huggingface.co/{model_id}/resolve/{revision}/model.safetensors"
|
|
96
|
+
raw_metadata = await fetch_safetensors_metadata(client=client, url=url, headers=headers)
|
|
97
|
+
|
|
98
|
+
# NOTE: Given that at the moment for Sentence Transformers only the Transformers module is considered, we
|
|
99
|
+
# set the default component name to `0_Transformer` as defined in modules.json
|
|
100
|
+
if "config_sentence_transformers.json" in file_paths:
|
|
101
|
+
raw_metadata = {"0_Transformer": raw_metadata}
|
|
102
|
+
else:
|
|
103
|
+
# NOTE: If the model is a transformers model, then we simply set the component name to `Transformer`, to
|
|
104
|
+
# make sure that we provide the expected input to the `parse_safetensors_metadata`
|
|
105
|
+
raw_metadata = {"Transformer": raw_metadata}
|
|
106
|
+
|
|
107
|
+
metadata = parse_safetensors_metadata(raw_metadata=raw_metadata)
|
|
108
|
+
elif "model.safetensors.index.json" in file_paths:
|
|
109
|
+
# TODO: We could eventually skip this request in favour of a greedy approach on trying to pull all the
|
|
110
|
+
# files following the formatting `model-00000-of-00000.safetensors`
|
|
111
|
+
url = f"https://huggingface.co/{model_id}/resolve/{revision}/model.safetensors.index.json"
|
|
112
|
+
files_index = await get_json_file(client=client, url=url, headers=headers)
|
|
113
|
+
|
|
114
|
+
urls = {
|
|
115
|
+
f"https://huggingface.co/{model_id}/resolve/{revision}/{f}"
|
|
116
|
+
for f in set(files_index["weight_map"].values())
|
|
117
|
+
}
|
|
118
|
+
|
|
119
|
+
semaphore = asyncio.Semaphore(MAX_CONCURRENCY)
|
|
120
|
+
|
|
121
|
+
async def fetch_with_semaphore(url: str) -> Dict[str, Any]:
|
|
122
|
+
async with semaphore:
|
|
123
|
+
return await fetch_safetensors_metadata(client=client, url=url, headers=headers)
|
|
124
|
+
|
|
125
|
+
tasks = [asyncio.create_task(fetch_with_semaphore(url)) for url in urls]
|
|
126
|
+
metadata_list: List[Dict[str, Any]] = await asyncio.gather(*tasks, return_exceptions=False)
|
|
127
|
+
|
|
128
|
+
raw_metadata = reduce(lambda acc, metadata: acc | metadata, metadata_list, {})
|
|
129
|
+
|
|
130
|
+
# NOTE: Given that at the moment for Sentence Transformers only the Transformers module is considered, we
|
|
131
|
+
# set the default component name to `0_Transformer` as defined in modules.json
|
|
132
|
+
if "config_sentence_transformers.json" in file_paths:
|
|
133
|
+
raw_metadata = {"0_Transformer": raw_metadata}
|
|
134
|
+
else:
|
|
135
|
+
# NOTE: If the model is a transformers model, then we simply set the component name to `Transformer`, to
|
|
136
|
+
# make sure that we provide the expected input to the `parse_safetensors_metadata`
|
|
137
|
+
raw_metadata = {"Transformer": raw_metadata}
|
|
138
|
+
|
|
139
|
+
metadata = parse_safetensors_metadata(raw_metadata=raw_metadata)
|
|
140
|
+
elif "model_index.json" in file_paths:
|
|
141
|
+
url = f"https://huggingface.co/{model_id}/resolve/{revision}/model_index.json"
|
|
142
|
+
files_index = await get_json_file(client=client, url=url, headers=headers)
|
|
143
|
+
paths = {k for k, _ in files_index.items() if not k.startswith("_")}
|
|
144
|
+
|
|
145
|
+
path_urls: Dict[str, List[str]] = {}
|
|
146
|
+
for path in paths:
|
|
147
|
+
if f"{path}/diffusion_pytorch_model.safetensors" in file_paths:
|
|
148
|
+
path_urls[path] = [
|
|
149
|
+
f"https://huggingface.co/{model_id}/resolve/{revision}/{path}/diffusion_pytorch_model.safetensors"
|
|
150
|
+
]
|
|
151
|
+
elif f"{path}/model.safetensors" in file_paths:
|
|
152
|
+
path_urls[path] = [
|
|
153
|
+
f"https://huggingface.co/{model_id}/resolve/{revision}/{path}/model.safetensors"
|
|
154
|
+
]
|
|
155
|
+
elif f"{path}/diffusion_pytorch_model.safetensors.index.json" in file_paths:
|
|
156
|
+
url = f"https://huggingface.co/{model_id}/resolve/{revision}/{path}/diffusion_pytorch_model.safetensors.index.json"
|
|
157
|
+
files_index = await get_json_file(client=client, url=url, headers=headers)
|
|
158
|
+
path_urls[path] = [
|
|
159
|
+
f"https://huggingface.co/{model_id}/resolve/{revision}/{path}/{f}"
|
|
160
|
+
for f in set(files_index["weight_map"].values())
|
|
161
|
+
]
|
|
162
|
+
elif f"{path}/model.safetensors.index.json" in file_paths:
|
|
163
|
+
url = (
|
|
164
|
+
f"https://huggingface.co/{model_id}/resolve/{revision}/{path}/model.safetensors.index.json"
|
|
165
|
+
)
|
|
166
|
+
files_index = await get_json_file(client=client, url=url, headers=headers)
|
|
167
|
+
path_urls[path] = [
|
|
168
|
+
f"https://huggingface.co/{model_id}/resolve/{revision}/{path}/{f}"
|
|
169
|
+
for f in set(files_index["weight_map"].values())
|
|
170
|
+
]
|
|
171
|
+
|
|
172
|
+
semaphore = asyncio.Semaphore(MAX_CONCURRENCY)
|
|
173
|
+
|
|
174
|
+
async def fetch_with_semaphore(url: str) -> Dict[str, Any]:
|
|
175
|
+
async with semaphore:
|
|
176
|
+
return await fetch_safetensors_metadata(client=client, url=url, headers=headers)
|
|
177
|
+
|
|
178
|
+
raw_metadata = {}
|
|
179
|
+
for path, urls in path_urls.items():
|
|
180
|
+
tasks = [asyncio.create_task(fetch_with_semaphore(url)) for url in urls]
|
|
181
|
+
metadata_list: List[Dict[str, Any]] = await asyncio.gather(*tasks, return_exceptions=False)
|
|
182
|
+
raw_metadata[path] = reduce(lambda acc, metadata: acc | metadata, metadata_list, {})
|
|
183
|
+
|
|
184
|
+
metadata = parse_safetensors_metadata(raw_metadata=raw_metadata)
|
|
185
|
+
else:
|
|
186
|
+
raise RuntimeError(
|
|
187
|
+
"NONE OF `model.safetensors`, `model.safetensors.index.json`, `model_index.json` HAS BEEN FOUND"
|
|
188
|
+
)
|
|
189
|
+
|
|
190
|
+
if json_output:
|
|
191
|
+
out = {"model_id": model_id, "revision": revision, **asdict(metadata)}
|
|
192
|
+
print(json.dumps(out))
|
|
193
|
+
else:
|
|
194
|
+
print_report(
|
|
195
|
+
model_id=model_id,
|
|
196
|
+
revision=revision,
|
|
197
|
+
metadata=metadata,
|
|
198
|
+
ignore_table_width=ignore_table_width,
|
|
199
|
+
)
|
|
200
|
+
|
|
201
|
+
|
|
202
|
+
def main() -> None:
|
|
203
|
+
parser = argparse.ArgumentParser()
|
|
204
|
+
parser.add_argument("--model-id", required=True, help="Model ID on the Hugging Face Hub")
|
|
205
|
+
parser.add_argument(
|
|
206
|
+
"--revision",
|
|
207
|
+
default="main",
|
|
208
|
+
help="Model revision on the Hugging Face Hub",
|
|
209
|
+
)
|
|
210
|
+
parser.add_argument(
|
|
211
|
+
"--json-output",
|
|
212
|
+
action="store_true",
|
|
213
|
+
help="Whether to provide the output as a JSON instead of printed as table.",
|
|
214
|
+
)
|
|
215
|
+
parser.add_argument(
|
|
216
|
+
"--ignore-table-width",
|
|
217
|
+
action="store_true",
|
|
218
|
+
help="Whether to ignore the maximum recommended table width, in case the `--model-id` and/or `--revision` cause a row overflow when printing those.",
|
|
219
|
+
)
|
|
220
|
+
|
|
221
|
+
args = parser.parse_args()
|
|
222
|
+
|
|
223
|
+
asyncio.run(
|
|
224
|
+
run(
|
|
225
|
+
model_id=args.model_id,
|
|
226
|
+
revision=args.revision,
|
|
227
|
+
# NOTE: Below are the arguments that affect the output format
|
|
228
|
+
json_output=args.json_output,
|
|
229
|
+
ignore_table_width=args.ignore_table_width,
|
|
230
|
+
)
|
|
231
|
+
)
|
hf_mem/metadata.py
ADDED
|
@@ -0,0 +1,61 @@
|
|
|
1
|
+
import math
|
|
2
|
+
from dataclasses import dataclass
|
|
3
|
+
from typing import Any, Dict
|
|
4
|
+
|
|
5
|
+
from hf_mem.types import SafetensorsDtypes, get_safetensors_dtype_bytes
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
@dataclass
|
|
9
|
+
class DtypeMetadata:
|
|
10
|
+
param_count: int
|
|
11
|
+
bytes_count: int
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
@dataclass
|
|
15
|
+
class ComponentMetadata:
|
|
16
|
+
dtypes: Dict[SafetensorsDtypes, DtypeMetadata]
|
|
17
|
+
param_count: int
|
|
18
|
+
bytes_count: int
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
@dataclass
|
|
22
|
+
class SafetensorsMetadata:
|
|
23
|
+
components: Dict[str, ComponentMetadata]
|
|
24
|
+
param_count: int
|
|
25
|
+
bytes_count: int
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def parse_safetensors_metadata(
|
|
29
|
+
raw_metadata: Dict[str, Dict[str, Any]],
|
|
30
|
+
) -> SafetensorsMetadata:
|
|
31
|
+
components = {}
|
|
32
|
+
total_param_count, total_bytes_count = 0, 0
|
|
33
|
+
|
|
34
|
+
for name, metadata in raw_metadata.items():
|
|
35
|
+
component = ComponentMetadata(dtypes={}, param_count=0, bytes_count=0)
|
|
36
|
+
for key, value in metadata.items():
|
|
37
|
+
if key in {"__metadata__"}:
|
|
38
|
+
continue
|
|
39
|
+
|
|
40
|
+
dtype = value["dtype"]
|
|
41
|
+
if dtype not in component.dtypes:
|
|
42
|
+
component.dtypes[dtype] = DtypeMetadata(param_count=0, bytes_count=0)
|
|
43
|
+
|
|
44
|
+
dtype_bytes = get_safetensors_dtype_bytes(dtype)
|
|
45
|
+
current_shape = math.prod(value["shape"])
|
|
46
|
+
current_shape_bytes = current_shape * dtype_bytes
|
|
47
|
+
|
|
48
|
+
component.dtypes[dtype].param_count += current_shape
|
|
49
|
+
component.dtypes[dtype].bytes_count += current_shape_bytes
|
|
50
|
+
component.param_count += current_shape
|
|
51
|
+
component.bytes_count += current_shape_bytes
|
|
52
|
+
total_param_count += current_shape
|
|
53
|
+
total_bytes_count += current_shape_bytes
|
|
54
|
+
|
|
55
|
+
components[name] = component
|
|
56
|
+
|
|
57
|
+
return SafetensorsMetadata(
|
|
58
|
+
components=components,
|
|
59
|
+
param_count=total_param_count,
|
|
60
|
+
bytes_count=total_bytes_count,
|
|
61
|
+
)
|
hf_mem/print.py
ADDED
|
@@ -0,0 +1,196 @@
|
|
|
1
|
+
import warnings
|
|
2
|
+
from typing import Literal, Optional
|
|
3
|
+
|
|
4
|
+
from hf_mem.metadata import SafetensorsMetadata
|
|
5
|
+
|
|
6
|
+
MIN_NAME_LEN = 5
|
|
7
|
+
MAX_NAME_LEN = 14
|
|
8
|
+
MIN_DATA_LEN = 20
|
|
9
|
+
MAX_DATA_LEN = 64
|
|
10
|
+
BORDERS_AND_PADDING = 7
|
|
11
|
+
|
|
12
|
+
BOX = {
|
|
13
|
+
"tl": "┌",
|
|
14
|
+
"tr": "┐",
|
|
15
|
+
"bl": "└",
|
|
16
|
+
"br": "┘",
|
|
17
|
+
"ht": "─",
|
|
18
|
+
"vt": "│",
|
|
19
|
+
"tsep": "┬",
|
|
20
|
+
"bsep": "┴",
|
|
21
|
+
"lm": "├",
|
|
22
|
+
"rm": "┤",
|
|
23
|
+
"mm": "┼",
|
|
24
|
+
}
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def _print_with_color(content: str) -> None:
|
|
28
|
+
print(f"\x1b[38;2;244;183;63m{content}\x1b[0m")
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def _print_header(current_len: int) -> None:
|
|
32
|
+
length = current_len + MAX_NAME_LEN + BORDERS_AND_PADDING
|
|
33
|
+
top = BOX["tl"] + (BOX["tsep"] * (length - 2)) + BOX["tr"]
|
|
34
|
+
_print_with_color(top)
|
|
35
|
+
|
|
36
|
+
bottom = BOX["lm"] + (BOX["bsep"] * (length - 2)) + BOX["rm"]
|
|
37
|
+
_print_with_color(bottom)
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
def _print_centered(text: str, current_len: int) -> None:
|
|
41
|
+
max_len = current_len + MAX_NAME_LEN - BORDERS_AND_PADDING
|
|
42
|
+
total_width = max_len + 12
|
|
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: Optional[Literal["top", "top-continue", "bottom", "bottom-continue"]] = None,
|
|
52
|
+
) -> None:
|
|
53
|
+
match side:
|
|
54
|
+
case "top":
|
|
55
|
+
left, mid, right = BOX["lm"], BOX["tsep"], BOX["rm"]
|
|
56
|
+
case "top-continue":
|
|
57
|
+
left, mid, right = BOX["lm"], BOX["bsep"], BOX["rm"]
|
|
58
|
+
case "bottom":
|
|
59
|
+
left, mid, right = BOX["bl"], BOX["bsep"], BOX["br"]
|
|
60
|
+
case "bottom-continue":
|
|
61
|
+
left, mid, right = BOX["lm"], BOX["bsep"], BOX["rm"]
|
|
62
|
+
case _:
|
|
63
|
+
left, mid, right = BOX["lm"], BOX["mm"], BOX["rm"]
|
|
64
|
+
|
|
65
|
+
name_col_inner = MAX_NAME_LEN + 2
|
|
66
|
+
data_col_inner = current_len + 1
|
|
67
|
+
|
|
68
|
+
line = left
|
|
69
|
+
line += BOX["ht"] * name_col_inner
|
|
70
|
+
line += mid
|
|
71
|
+
line += BOX["ht"] * data_col_inner
|
|
72
|
+
line += right
|
|
73
|
+
_print_with_color(line)
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
def _format_name(name: str) -> str:
|
|
77
|
+
if len(name) < MIN_NAME_LEN:
|
|
78
|
+
return f"{name:<{MIN_NAME_LEN}}"
|
|
79
|
+
if len(name) > MAX_NAME_LEN:
|
|
80
|
+
return name[: MAX_NAME_LEN - 3] + "..."
|
|
81
|
+
return f"{name:<{MAX_NAME_LEN}}"
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
def _print_row(name: str, text: str, current_len: int) -> None:
|
|
85
|
+
name_fmt = _format_name(name)
|
|
86
|
+
data_fmt = f"{str(text):<{current_len}}"
|
|
87
|
+
_print_with_color(f"{BOX['vt']} {name_fmt} {BOX['vt']} {data_fmt} {BOX['vt']}")
|
|
88
|
+
|
|
89
|
+
|
|
90
|
+
def _make_bar(used: float, total: float, width: int) -> str:
|
|
91
|
+
if total <= 0:
|
|
92
|
+
return "░" * width
|
|
93
|
+
frac = min(max(used / total, 0.0), 1.0)
|
|
94
|
+
filled = int(round(frac * width))
|
|
95
|
+
filled = max(0, min(width, filled))
|
|
96
|
+
return "█" * filled + "░" * (width - filled)
|
|
97
|
+
|
|
98
|
+
|
|
99
|
+
def _format_short_number(n: float) -> str:
|
|
100
|
+
n = float(n)
|
|
101
|
+
for unit in ("", "K", "M", "B", "T"):
|
|
102
|
+
if abs(n) < 1000.0:
|
|
103
|
+
return f"{int(n)}" if unit == "" else f"{n:.2f}{unit}"
|
|
104
|
+
n /= 1000.0
|
|
105
|
+
return f"{n:.2f}P"
|
|
106
|
+
|
|
107
|
+
|
|
108
|
+
def _bytes_to_gb(nbytes: int) -> float:
|
|
109
|
+
return nbytes / (1024**3)
|
|
110
|
+
|
|
111
|
+
|
|
112
|
+
def print_report(
|
|
113
|
+
model_id: str,
|
|
114
|
+
revision: str,
|
|
115
|
+
metadata: SafetensorsMetadata,
|
|
116
|
+
ignore_table_width: bool = False,
|
|
117
|
+
) -> None:
|
|
118
|
+
rows = [
|
|
119
|
+
"INFERENCE MEMORY ESTIMATE FOR",
|
|
120
|
+
f"https://hf.co/{model_id} @ {revision}",
|
|
121
|
+
"TOTAL MEMORY",
|
|
122
|
+
"REQUIREMENTS",
|
|
123
|
+
]
|
|
124
|
+
|
|
125
|
+
for name, nested_metadata in metadata.components.items():
|
|
126
|
+
if len(metadata.components) > 1:
|
|
127
|
+
rows.append(name)
|
|
128
|
+
|
|
129
|
+
for dtype, dtype_metadata in nested_metadata.dtypes.items():
|
|
130
|
+
rows.append(f"{dtype} {dtype_metadata.param_count} {dtype_metadata.bytes_count}")
|
|
131
|
+
|
|
132
|
+
max_len = max(len(str(r)) for r in rows)
|
|
133
|
+
|
|
134
|
+
if max_len > MAX_DATA_LEN and ignore_table_width is False:
|
|
135
|
+
warnings.warn(
|
|
136
|
+
f"Given that the provided `--model-id {model_id}` (with `--revision {revision}`) is longer than {MAX_DATA_LEN} characters, the table width will be expanded to fit the provided values within their row, but it might lead to unexpected table views. If you'd like to ignore the limit, then provide the `--ignore-table-width` flag to ignore the {MAX_DATA_LEN} width limit, to simply accommodate to whatever the longest text length is."
|
|
137
|
+
)
|
|
138
|
+
|
|
139
|
+
current_len = min(max_len, MAX_DATA_LEN) if ignore_table_width is False else max_len
|
|
140
|
+
|
|
141
|
+
_print_header(current_len)
|
|
142
|
+
_print_centered("INFERENCE MEMORY ESTIMATE FOR", current_len)
|
|
143
|
+
_print_centered(f"https://hf.co/{model_id} @ {revision}", current_len)
|
|
144
|
+
_print_divider(current_len + 1, "top")
|
|
145
|
+
|
|
146
|
+
total_text = (
|
|
147
|
+
f"{_bytes_to_gb(metadata.bytes_count):.2f} GB ({_format_short_number(metadata.param_count)} params)"
|
|
148
|
+
)
|
|
149
|
+
_print_row("TOTAL MEMORY", total_text, current_len)
|
|
150
|
+
|
|
151
|
+
total_bar = _make_bar(metadata.bytes_count, metadata.bytes_count, current_len)
|
|
152
|
+
_print_row("REQUIREMENTS", total_bar, current_len)
|
|
153
|
+
|
|
154
|
+
for key, value in metadata.components.items():
|
|
155
|
+
if len(metadata.components) > 1:
|
|
156
|
+
_print_divider(current_len + 1, "top-continue")
|
|
157
|
+
_print_centered(
|
|
158
|
+
f"{key.upper()} ({_bytes_to_gb(value.bytes_count):.2f} GB)",
|
|
159
|
+
current_len,
|
|
160
|
+
)
|
|
161
|
+
|
|
162
|
+
_print_divider(current_len + 1, "top")
|
|
163
|
+
else:
|
|
164
|
+
_print_divider(current_len + 1)
|
|
165
|
+
|
|
166
|
+
max_length = max(
|
|
167
|
+
[
|
|
168
|
+
len(f"{_format_short_number(dtype_metadata.param_count)} PARAMS")
|
|
169
|
+
for _, dtype_metadata in value.dtypes.items()
|
|
170
|
+
]
|
|
171
|
+
)
|
|
172
|
+
for idx, (dtype, dtype_metadata) in enumerate(value.dtypes.items()):
|
|
173
|
+
gb_text = (
|
|
174
|
+
f"{_bytes_to_gb(dtype_metadata.bytes_count):.2f} / {_bytes_to_gb(metadata.bytes_count):.2f} GB"
|
|
175
|
+
)
|
|
176
|
+
_print_row(
|
|
177
|
+
dtype.upper() + " " * (max_length - len(dtype)),
|
|
178
|
+
gb_text,
|
|
179
|
+
current_len,
|
|
180
|
+
)
|
|
181
|
+
|
|
182
|
+
bar = _make_bar(
|
|
183
|
+
_bytes_to_gb(dtype_metadata.bytes_count),
|
|
184
|
+
_bytes_to_gb(metadata.bytes_count),
|
|
185
|
+
current_len,
|
|
186
|
+
)
|
|
187
|
+
_print_row(
|
|
188
|
+
f"{_format_short_number(dtype_metadata.param_count)} PARAMS",
|
|
189
|
+
bar,
|
|
190
|
+
current_len,
|
|
191
|
+
)
|
|
192
|
+
|
|
193
|
+
if idx < len(value.dtypes) - 1:
|
|
194
|
+
_print_divider(current_len + 1)
|
|
195
|
+
|
|
196
|
+
_print_divider(current_len + 1, "bottom")
|
hf_mem/py.typed
ADDED
|
File without changes
|
hf_mem/types.py
ADDED
|
@@ -0,0 +1,32 @@
|
|
|
1
|
+
from typing import Literal
|
|
2
|
+
|
|
3
|
+
SafetensorsDtypes = Literal[
|
|
4
|
+
"F64",
|
|
5
|
+
"I64",
|
|
6
|
+
"U64",
|
|
7
|
+
"F32",
|
|
8
|
+
"I32",
|
|
9
|
+
"U32",
|
|
10
|
+
"F16",
|
|
11
|
+
"BF16",
|
|
12
|
+
"I16",
|
|
13
|
+
"U16",
|
|
14
|
+
"F8_E5M2",
|
|
15
|
+
"F8_E4M3",
|
|
16
|
+
"I8",
|
|
17
|
+
"U8",
|
|
18
|
+
]
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def get_safetensors_dtype_bytes(dtype: str) -> int:
|
|
22
|
+
match dtype:
|
|
23
|
+
case "F64" | "I64" | "U64":
|
|
24
|
+
return 8
|
|
25
|
+
case "F32" | "I32" | "U32":
|
|
26
|
+
return 4
|
|
27
|
+
case "F16" | "BF16" | "I16" | "U16":
|
|
28
|
+
return 2
|
|
29
|
+
case "F8_E5M2" | "F8_E4M3" | "I8" | "U8":
|
|
30
|
+
return 1
|
|
31
|
+
case _:
|
|
32
|
+
raise RuntimeError(f"DTYPE={dtype} NOT HANDLED")
|
|
@@ -0,0 +1,40 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: hf-mem
|
|
3
|
+
Version: 0.2.2
|
|
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 <36760800+alvarobartt@users.noreply.github.com>
|
|
7
|
+
License-Expression: MIT
|
|
8
|
+
License-File: LICENSE
|
|
9
|
+
Requires-Dist: httpx[http2]>=0.28.1
|
|
10
|
+
Requires-Python: >=3.10
|
|
11
|
+
Description-Content-Type: text/markdown
|
|
12
|
+
|
|
13
|
+
<img src="https://github.com/user-attachments/assets/509a8244-8a91-4051-b337-41b7b2fe0e2f" />
|
|
14
|
+
|
|
15
|
+
---
|
|
16
|
+
|
|
17
|
+
`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.
|
|
18
|
+
|
|
19
|
+
`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.
|
|
20
|
+
|
|
21
|
+
Read more information about `hf-mem` in [this short-form post](https://alvarobartt.com/hf-mem).
|
|
22
|
+
|
|
23
|
+
## Usage
|
|
24
|
+
|
|
25
|
+
```bash
|
|
26
|
+
uvx hf-mem --model-id MiniMaxAI/MiniMax-M2
|
|
27
|
+
```
|
|
28
|
+
|
|
29
|
+
<img src="https://github.com/user-attachments/assets/530f8b14-a415-4fd6-9054-bcd81cafae09" />
|
|
30
|
+
|
|
31
|
+
```bash
|
|
32
|
+
uvx hf-mem --model-id Qwen/Qwen-Image
|
|
33
|
+
```
|
|
34
|
+
|
|
35
|
+
<img src="https://github.com/user-attachments/assets/cd4234ec-bdcc-4db4-8b01-0ac9b5cd390c" />
|
|
36
|
+
|
|
37
|
+
## References
|
|
38
|
+
|
|
39
|
+
- [Safetensors Metadata parsing](https://huggingface.co/docs/safetensors/en/metadata_parsing)
|
|
40
|
+
- [usgraphics - TR-100 Machine Report](https://github.com/usgraphics/usgc-machine-report)
|
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
hf_mem/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
2
|
+
hf_mem/cli.py,sha256=SV6HvN1emkm_fWoEq_V14RNXUwer0Y5wTmVqOtsc5Ns,10233
|
|
3
|
+
hf_mem/metadata.py,sha256=yWamwZg-zc7luhvZx8tqIe8cjA9o44qrboi8DZ83lIg,1769
|
|
4
|
+
hf_mem/print.py,sha256=TJvDFPzhYVDpjgyw4yxU91BuEb9k6x_5WdkIDW9p1_8,6365
|
|
5
|
+
hf_mem/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
6
|
+
hf_mem/types.py,sha256=GLs8S3lIP9JkwdoDc8cwHn3IlBOAIbH-NihtzDsmzyI,619
|
|
7
|
+
hf_mem-0.2.2.dist-info/licenses/LICENSE,sha256=TUO4cSlmwlYsHZfepsmo1zLoYhmT0RImXoRihVYNHhw,1060
|
|
8
|
+
hf_mem-0.2.2.dist-info/WHEEL,sha256=KSLUh82mDPEPk0Bx0ScXlWL64bc8KmzIPNcpQZFV-6E,79
|
|
9
|
+
hf_mem-0.2.2.dist-info/entry_points.txt,sha256=p6o2afl5zELLNtkKwQxfLdOpal2nRlcshDZSfV02pOY,44
|
|
10
|
+
hf_mem-0.2.2.dist-info/METADATA,sha256=AF2JzRasS5OPzualhbHkPvxmeHo6w7AKN6l0bOFpYQc,1647
|
|
11
|
+
hf_mem-0.2.2.dist-info/RECORD,,
|
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
Copyright (c) 2025 Alvaro Bartolome
|
|
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.
|